Add type hints for generate

jncraton · Jan 4, 2024 · 7254bd7 · 7254bd7
1 parent 1a065a3
commit 7254bd7
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/languagemodels/inference.py b/languagemodels/inference.py
@@ -1,3 +1,4 @@
+from typing import List
 import requests
 import re
 import os
@@ -111,14 +112,14 @@ def chat_oa(engine, prompt, max_tokens=200, temperature=0):
 
 
 def generate(
-    instructions,
-    max_tokens=200,
-    temperature=0.1,
-    topk=1,
-    repetition_penalty=1.3,
-    prefix="",
-    suppress=[],
-    model="instruct",
+    instructions: List[str],
+    max_tokens: int = 200,
+    temperature: float = 0.1,
+    topk: int = 1,
+    repetition_penalty: float = 1.3,
+    prefix: str = "",
+    suppress: List[str] = [],
+    model: str = "instruct",
 ):
     """Generates completions for a prompt