From f7da75d18920e165d0f53c009618145a58c848ae Mon Sep 17 00:00:00 2001 From: Prikshit7766 Date: Wed, 18 Oct 2023 14:11:02 +0530 Subject: [PATCH] fix params for PretrainedModelForSensitivityTest --- langtest/modelhandler/llm_modelhandler.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/langtest/modelhandler/llm_modelhandler.py b/langtest/modelhandler/llm_modelhandler.py index d9f45baae..4cc1287f9 100644 --- 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