import os
import subprocess

# Uninstall the existing llama-cpp-python package
subprocess.run("pip uninstall -y llama-cpp-python", shell=True)

# Get the path to the CUDA compiler
cuda_path = os.environ.get("CUDA_PATH", "/usr/local/cuda")
nvcc_path = os.path.join(cuda_path, "bin", "nvcc")

# Set the CMAKE_CUDA_COMPILER environment variable
os.environ["CMAKE_CUDA_COMPILER"] = nvcc_path

# Install llama-cpp-python with CUDA support
install_command = "pip install llama-cpp-python --install-option='--cmake-args=-DLLAMA_CUDA=on'"
subprocess.run(install_command, shell=True)