fixed env variable issues

qcri · Sep 17, 2023 · 575900e · 575900e
1 parent 5f60157
commit 575900e
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 8 deletions.
diff --git a/...r/MT/AraBench_ara2en_FastChat_ZeroShot.py → ...ar/MT/AraBench_ar2en_FastChat_ZeroShot.py b/...r/MT/AraBench_ara2en_FastChat_ZeroShot.py → ...ar/MT/AraBench_ar2en_FastChat_ZeroShot.py
@@ -60,10 +60,6 @@ def config():
                     "task_args": {},
                     "model": FastChatModel,
                     "model_args": {
-                        "api_type": "openai",
-                        "api_key": os.environ["OPENAI_API_KEY"],
-                        "api_base": os.environ["OPENAI_API_URL"],
-                        "model_name": os.environ["OPENAI_MODEL"],
                         "max_tries": 5,
                     },
                     "general_args": {"data_path": "data/MT/"},

diff --git a/llmebench/models/FastChat.py b/llmebench/models/FastChat.py
@@ -4,10 +4,22 @@
 
 
 class FastChatModel(OpenAIModel):
-    def __init__(self, api_base, api_key, model_name, **kwargs):
-        api_base = api_base or os.getenv("FASTCHAT_API_BASE")
-        api_key = api_key or os.getenv("FASTCHAT_API_KEY")
-        model_name = model_name or os.getenv("FASTCHAT_MODEL_NAME")
+    def __init__(self, api_base=None, api_key=None, model_name=None, **kwargs):
+        self.api_base = api_base or os.getenv("OPENAI_API_BASE")
+        self.api_key = api_key or os.getenv("OPENAI_API_KEY")
+        self.model_name = model_name or os.getenv("OPENAI_MODEL")
+        if self.api_base is None:
+            raise Exception(
+                "API url must be provided as model config or environment variable (`OPENAI_API_BASE`)"
+            )
+        if self.api_key is None:
+            raise Exception(
+                "API url must be provided as model config or environment variable (`OPENAI_API_KEY`)"
+            )
+        if self.model_name is None:
+            raise Exception(
+                "API url must be provided as model config or environment variable (`OPENAI_MODEL`)"
+            )
         # checks for valid config settings)
         super(FastChatModel, self).__init__(
             api_base=api_base, api_key=api_key, model_name=model_name, **kwargs