File size: 1,476 Bytes
053eae3 1af9d11 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
import subprocess
import re
import sys
def install_llama_cpp_python(cuda=None):
try:
# Detect CUDA version
result = subprocess.run(['nvcc', '--version'], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
version_output = result.stdout
version_match = re.search(r"release (\d+\.\d+),", version_output)
if not version_match:
raise RuntimeError("CUDA version not found in nvcc output.")
cuda_version = version_match.group(1)
version_map = {
"12.1": "cu121",
"12.2": "cu122",
"12.3": "cu123",
"12.4": "cu124",
}
major_minor = '.'.join(cuda_version.split('.')[:2])
if cuda is None:
cuda_suffix = version_map.get(major_minor)
else:
cuda_suffix = cuda
if not cuda_suffix:
raise ValueError(f"No suitable wheel found for CUDA version {cuda_version}.")
extra_index_url = f"https://abetlen.github.io/llama-cpp-python/whl/{cuda_suffix}"
subprocess.run([sys.executable, '-m', 'pip', 'install', 'llama-cpp-python', '--extra-index-url', extra_index_url])
print(f"Successfully installed llama-cpp-python with CUDA {cuda_suffix}")
except FileNotFoundError:
print("Error: nvcc (CUDA) is not installed or not in PATH.")
except Exception as e:
print(f"An error occurred: {e}")
# Run the function
install_llama_cpp_python('cu124') |