Handle non-speculative fallback

armbues · Jan 29, 2025 · 375e65c · 375e65c
1 parent eb407a2
commit 375e65c
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/sillm/experimental/speculative.py b/sillm/experimental/speculative.py
@@ -159,6 +159,7 @@ def from_model(llm: LLM):
         return SpeculativeEdit(llm)
 
     def __init__(self, model: LLM):
+        self.target_model = model
         self.model = model.model
         self.tokenizer = model.tokenizer
         self.args = model.args
@@ -203,6 +204,14 @@ def generate(self,
             "token_ids": []
         }
 
+        # Handle non-speculative generation
+        if lookahead <= 1:
+            for s, metadata in self.target_model.generate(prompt, temperature=temperature, max_tokens=max_tokens, token_ids=token_ids):
+                metadata["speculative"] = speculative
+                yield s, metadata
+
+            return
+
         # Define stop tokens
         stop_tokens = self.tokenizer.special_ids