CUDA 12 release (#140)

.github/workflows/build.yaml

@@ -41,7 +41,7 @@ jobs:
       matrix:
         os: [ubuntu-20.04, windows-latest]
         pyver: ["3.8", "3.9", "3.10", "3.11"]
-        cuda: ["11.8.0", "12.1.1"]
+        cuda: ["12.1.1"]
     defaults:
       run:
         shell: pwsh
@@ -105,6 +105,7 @@ jobs:
         run: |
           $env:CUDA_PATH = $env:CONDA_PREFIX
           $env:CUDA_HOME = $env:CONDA_PREFIX
+          $env:PYPI_BUILD = 1
 
           python setup.py sdist bdist_wheel
       

scripts/download_wheels.sh

 #!/bin/bash
 
-# Set the GitHub release URL
-RELEASE_URL="https://api.github.com/repos/casper-hansen/AutoAWQ/releases/tags/v0.0.1"
+# Set variables
+AWQ_VERSION="0.1.6"
+CUDA_VERSION="cu1180"
+RELEASE_URL="https://api.github.com/repos/casper-hansen/AutoAWQ/releases/tags/v${AWQ_VERSION}"
 
 # Create a directory to download the wheels
 mkdir -p dist
@@ -10,12 +12,21 @@ cd dist
 # Download all the wheel files from the GitHub release
 curl -s $RELEASE_URL | \
     jq -r ".assets[].browser_download_url" | \
-    grep .whl | \
+    grep '\.whl' | \
+    grep "$CUDA_VERSION" | \
     xargs -n 1 wget
 
 # Rename the wheels from 'linux_x86_64' to 'manylinux_x86_64'
-for file in *linux_x86_64.whl; do
-    mv "$file" "$(echo $file | sed 's/linux_x86_64/manylinux2014_x86_64/')"
+# Remove CUDA version from filename
+for file in *"$CUDA_VERSION"*.whl; do
+    # First, rename the platform from 'linux_x86_64' to 'manylinux2014_x86_64'
+    intermediate_name=$(echo "$file" | sed 's/linux_x86_64/manylinux2014_x86_64/')
+    
+    # Then, remove the CUDA version from the filename
+    newname=$(echo "$intermediate_name" | sed "s/+${CUDA_VERSION}//")
+    
+    mv -v "$file" "$newname"
 done
 
+
 cd ..

setup.py

@@ -7,14 +7,18 @@ from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CUDAExtension
 
 os.environ["CC"] = "g++"
 os.environ["CXX"] = "g++"
-
-try:
-    CUDA_VERSION = "".join(os.environ.get("CUDA_VERSION", torch.version.cuda).split("."))
-except Exception as ex:
+AUTOAWQ_VERSION = "0.1.6"
+PYPI_BUILD = os.getenv("PYPI_BUILD", "0") == "1"
+
+if not PYPI_BUILD:
+    try:
+        CUDA_VERSION = "".join(os.environ.get("CUDA_VERSION", torch.version.cuda).split("."))[:3]
+        AUTOAWQ_VERSION += f"cu+{CUDA_VERSION}"
+    except Exception as ex:
         raise RuntimeError("Your system must have an Nvidia GPU for installing AutoAWQ")
 
 common_setup_kwargs = {
-    "version": f"0.1.6+cu{CUDA_VERSION}",
+    "version": AUTOAWQ_VERSION,
     "name": "autoawq",
     "author": "Casper Hansen",
     "license": "MIT",
@@ -39,7 +43,7 @@ common_setup_kwargs = {
 }
 
 requirements = [
-    "torch>=2.0.0",
+    "torch>=2.1.0",
     "transformers>=4.34.0",
     "tokenizers>=0.12.1",
     "accelerate",