Update vllm.py

BatsResearch · Mar 4, 2024 · bd6aab9 · bd6aab9
1 parent 3c2f3f0
commit bd6aab9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/alfred/fm/vllm.py b/alfred/fm/vllm.py
@@ -22,7 +22,7 @@ class vLLMModel(LocalAccessFoundationModel):
     source: https://github.com/vllm-project/vllm
     """
 
-    def __init__(self, model: str, local_dir: str = None, **kwargs: Any):
+    def __init__(self, model: str, local_path: str = None, **kwargs: Any):
         """
         Initialize a VLLM with MultiGPU.
 
@@ -33,7 +33,7 @@ def __init__(self, model: str, local_dir: str = None, **kwargs: Any):
         super().__init__(model)
         self.gpu_count = torch.cuda.device_count()
         self.model = LLM(
-            local_dir if local_dir is not None else model,
+            local_path if local_path is not None else model,
             tensor_parallel_size=self.gpu_count,
         )