Update vllm.py (#134)

zou-group · Oct 24, 2024 · 2042663 · 2042663
1 parent adadbd2
commit 2042663
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/textgrad/engine/vllm.py b/textgrad/engine/vllm.py
@@ -18,14 +18,15 @@ def __init__(
         self,
         model_string="meta-llama/Meta-Llama-3-8B-Instruct",
         system_prompt=DEFAULT_SYSTEM_PROMPT,
+        **llm_config,
     ):
         root = platformdirs.user_cache_dir("textgrad")
         cache_path = os.path.join(root, f"cache_vllm_{model_string}.db")
         super().__init__(cache_path=cache_path)
 
         self.model_string = model_string
         self.system_prompt = system_prompt
-        self.client = LLM(self.model_string)
+        self.client = LLM(self.model_string, **llm_config)
         self.tokenizer = self.client.get_tokenizer()
 
     def generate(