Update inference.py (#7)

- Update inference.py (c08633e8693e1a0b1b3b853b01a6df653192d516)


Co-authored-by: GRINDA AI <[email protected]>

inference.py

@@ -5,12 +5,19 @@ import sys
 import torch
 from typing import List, Dict
 
-# Ensure vllm is installed
+# Ensure vllm is installed and specify version to match CUDA compatibility
 try:
  import vllm
 except ImportError:
- subprocess.check_call([sys.executable, "-m", "pip", "install", "vllm"])
+ # Check CUDA version and install the correct vllm version
+ cuda_version = torch.version.cuda
+ if cuda_version == "11.8":
+ vllm_version = "v0.6.1.post1"
+ pip_cmd = f"pip install https://github.com/vllm-project/vllm/releases/download/{vllm_version}/vllm-{vllm_version}+cu118-cp310-cp310-manylinux1_x86_64.whl --extra-index-url https://download.pytorch.org/whl/cu118"
+ else:
+ raise RuntimeError(f"Unsupported CUDA version: {cuda_version}")
 
+ subprocess.check_call([sys.executable, "-m", "pip", "install", pip_cmd])
 # Import the necessary modules after installation
 from vllm import LLM, SamplingParams
 from vllm.utils import random_uuid