Update generator.py to allow completion_type to be passed as kwarg

McGill-NLP · May 29, 2024 · 3c65fa4 · 3c65fa4
1 parent 5a219c0
commit 3c65fa4
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/instruct_qa/generation/generator.py b/instruct_qa/generation/generator.py
@@ -70,6 +70,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"]
+
         assert (
             self.model_name in self.model_map
         ), "You should add the model name to the model -> endpoint compatibility mappings."