Browse Source

Fix triton env marker

David Kim 2 months ago
parent
commit
00b71fd4fa
2 changed files with 5 additions and 5 deletions
  1. 1 0
      requirements.txt
  2. 4 5
      setup.py

+ 1 - 0
requirements.txt

@@ -4,3 +4,4 @@ torch
 tqdm
 more-itertools
 tiktoken
+triton>=2.0.0,<3;platform_machine=="x86_64" and sys_platform=="linux" or sys_platform=="linux2"

+ 4 - 5
setup.py

@@ -1,6 +1,6 @@
-import os
 import platform
 import sys
+from pathlib import Path
 
 import pkg_resources
 from setuptools import find_packages, setup
@@ -13,7 +13,7 @@ def read_version(fname="whisper/version.py"):
 
 requirements = []
 if sys.platform.startswith("linux") and platform.machine() == "x86_64":
-    requirements.append("triton==2.0.0")
+    requirements.append("triton>=2.0.0,<3")
 
 setup(
     name="openai-whisper",
@@ -28,11 +28,10 @@ setup(
     url="https://github.com/openai/whisper",
     license="MIT",
     packages=find_packages(exclude=["tests*"]),
-    install_requires=requirements
-    + [
+    install_requires=[
         str(r)
         for r in pkg_resources.parse_requirements(
-            open(os.path.join(os.path.dirname(__file__), "requirements.txt"))
+            Path(__file__).with_name("requirements.txt").open()
         )
     ],
     entry_points={