🐛 fall back for model size

Wytamma · Mar 15, 2024 · cf257eb · cf257eb
1 parent b50bece
commit cf257eb
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/write_the/llm.py b/write_the/llm.py
@@ -23,7 +23,14 @@ def __init__(self, prompt: PromptTemplate, temperature=0, model_name="gpt-3.5-tu
         self.prompt_size = self.number_of_tokens(prompt.template)
         self.temperature = temperature
         self.model_name = model_name
-        self.max_tokens = int(models[model_name]["context_window"])
+        try:
+            self.max_tokens = int(models[model_name]["context_window"])
+        except KeyError:
+            self.max_tokens = 4096
+            if model_name.startswith('gpt-4'):
+                self.max_tokens = 8192
+            elif model_name.startswith('gpt-3'):
+                self.max_tokens = 4096
 
     async def run(self, code, **kwargs):
         """