changed GemmaLLMInterface

interface.py

@@ -6,15 +6,13 @@ import torch
 from transformers import TextIteratorStreamer
 from threading import Thread
 
-
-
 class GemmaLLMInterface(CustomLLM):
     model: Any
     tokenizer: Any
     context_window: int = 8192
     num_output: int = 2048
     model_name: str = "gemma_2"
-
+    
     def _format_prompt(self, message: str) -> str:
         return (
             f"<start_of_turn>user\n{message}<end_of_turn>\n<start_of_turn>model\n"
@@ -27,48 +25,51 @@ class GemmaLLMInterface(CustomLLM):
             num_output=self.num_output,
             model_name=self.model_name,
         )
-
-    @llm_completion_callback()
-    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+    
+    def _prepare_generation(self, prompt: str) -> tuple:
         prompt = self._format_prompt(prompt)
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.model.to(device)
-
-        # Tokenize prompt and move inputs to the correct device
+        
         inputs = self.tokenizer(prompt, return_tensors="pt", add_special_tokens=True).to(device)
-
-        # Ensure the input doesn't exceed the maximum token length
         if inputs["input_ids"].shape[1] > self.context_window:
             inputs["input_ids"] = inputs["input_ids"][:, -self.context_window:]
-
-        # Create a streamer to handle token streaming
+        
         streamer = TextIteratorStreamer(self.tokenizer, timeout=None, skip_prompt=True, skip_special_tokens=True)
-
-        # Generate kwargs for the model
+        
         generate_kwargs = {
             "input_ids": inputs["input_ids"],
             "streamer": streamer,
             "max_new_tokens": self.num_output,
             "do_sample": True,
-            "top_p": 0.9,  # You can tweak these sampling params based on your needs
+            "top_p": 0.9,
             "top_k": 50,
             "temperature": 0.7,
             "num_beams": 1,
             "repetition_penalty": 1.1,
         }
-
-        # Launch the generation in a separate thread to stream the output
+        
+        return streamer, generate_kwargs
+           
+    @llm_completion_callback()
+    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+        streamer, generate_kwargs = self._prepare_generation(prompt)
+        
         t = Thread(target=self.model.generate, kwargs=generate_kwargs)
         t.start()
-
-        # Collect the streamed response token by token
+        
         response = ""
         for new_token in streamer:
             response += new_token
-            yield CompletionResponse(text=response)
+        
+        return CompletionResponse(text=response)
 
     @llm_completion_callback()
     def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
-        # Use the complete method to stream the output in real-time
-        for response in self.complete(prompt):
-            yield response
+        streamer, generate_kwargs = self._prepare_generation(prompt)
+        
+        t = Thread(target=self.model.generate, kwargs=generate_kwargs)
+        t.start()
+        
+        for new_token in streamer:
+            yield CompletionResponse(text=new_token)