Add option for array and list prompt input

armbues · Jan 17, 2025 · 8830a6d · 8830a6d
1 parent 9d837b8
commit 8830a6d
Showing 1 changed file with 10 additions and 3 deletions.
diff --git a/sillm/core/llm.py b/sillm/core/llm.py
@@ -445,7 +445,7 @@ def completion(self, *args, **kwargs) -> str:
 
 def generate(model,
              tokenizer: Tokenizer,
-             prompt: str,
+             prompt: str | list | mx.array,
              cache: KVCache = None,
              max_tokens: int = 2048,
              temperature: float = 0.0,
@@ -463,8 +463,15 @@ def generate(model,
              ):
     start = time.perf_counter()
 
-    # Tokenize inputs
-    inputs = mx.array(tokenizer.encode(prompt))
+    # Pre-process inputs
+    if isinstance(prompt, str):
+        inputs = mx.array(tokenizer.encode(prompt))
+    elif isinstance(prompt, list):
+        inputs = mx.array(prompt)
+    elif isinstance(prompt, mx.array):
+        inputs = prompt
+    else:
+        raise ValueError("Prompt must be a string, list of tokens, or MX array")
 
     # Initialize metadata
     timing = {