[fix] LLMModel interface

HelloBroBro · Feb 28, 2024 · 4c68eb4 · 4c68eb4
1 parent 162d1b6
commit 4c68eb4
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/promptbench/models/__init__.py b/promptbench/models/__init__.py
@@ -68,7 +68,7 @@ class LLMModel(object):
  def model_list():
  return SUPPORTED_MODELS
 
- def __init__(self, model, max_new_tokens=20, temperature=0, device="cuda", dtype="auto", model_dir=None, system_prompt=None, api_key=None):
+ def __init__(self, model: str, max_new_tokens: int=20, temperature: float=0.0, device: str="cuda", dtype: str="auto", model_dir: str=None, system_prompt: str=None, api_key:str =None):
  self.model_name = model
  self.model = self._create_model(max_new_tokens, temperature, device, dtype, model_dir, system_prompt, api_key)