Update generator.py

McGill-NLP · May 29, 2024 · 009df2d · 009df2d
1 parent 3c65fa4
commit 009df2d
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/instruct_qa/generation/generator.py b/instruct_qa/generation/generator.py
@@ -62,6 +62,8 @@ def post_process_response(self, response):
 
 class GPTx(BaseGenerator):
     def __init__(self, *args, **kwargs):
+        completion_type = kwargs.pop("completion_type", None)
+
         super().__init__(*args, **kwargs)
         openai.api_key = self.api_key
         self.model_map = {
@@ -70,8 +72,9 @@ def __init__(self, *args, **kwargs):
             "text-davinci-003": "completions",
             "text-davinci-002": "completions",
         }
-        if "completion_type" in kwargs:
-            self.model_map[self.model_name] = kwargs["completion_type"]
+
+        if completion_type is not None:
+            self.model_map[model_name] = completion_type
 
         assert (
             self.model_name in self.model_map