fix params for PretrainedModelForSensitivityTest

JohnSnowLabs · Oct 18, 2023 · f7da75d · f7da75d
1 parent 859f3a6
commit f7da75d
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/langtest/modelhandler/llm_modelhandler.py b/langtest/modelhandler/llm_modelhandler.py
@@ -242,7 +242,7 @@ def __init__(self, model: str):
         self.model, self.embeddings_model = model
 
     @classmethod
-    def load_model(cls, path: str) -> tuple:
+    def load_model(cls, path: str, *args, **kwargs) -> tuple:
         """
         Load the pretrained language model and embeddings model from a given path.
 
@@ -262,6 +262,8 @@ def load_model(cls, path: str) -> tuple:
                 model_name=path,
                 temperature=0,
                 openai_api_key=os.environ["OPENAI_API_KEY"],
+                *args,
+                **kwargs,
             )
             embeddings_model = OpenAIEmbeddings(model="text-embedding-ada-002")
             return llm, embeddings_model