changed class interface

backend.py

@@ -28,19 +28,14 @@ model = AutoModelForCausalLM.from_pretrained(
     token=True)
 
 model.tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it")
-
 model.eval()
 
-#from accelerate import disk_offload
-#disk_offload(model=model, offload_dir="offload")
-
 # what models will be used by LlamaIndex:
 Settings.embed_model = InstructorEmbedding(model_name="hkunlp/instructor-base")
 
 Settings.llm = GemmaLLMInterface(model=model)
 #Settings.llm = GemmaLLMInterface(model_name=model_id)
 
-
 ############################---------------------------------
 
 # Get the parser

interface.py

@@ -36,17 +36,22 @@ class GemmaLLMInterface(CustomLLM):
         outputs = self.model.generate(**inputs, max_length=self.num_output)
         response = self.model.tokenizer.decode(outputs[0], skip_special_tokens=True)
         response = response[len(prompt):].strip()
-        # Ensure we always return a non-empty response
         return CompletionResponse(text=response if response else "No response generated.")
 
     @llm_completion_callback()
     def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
-        full_response = self.complete(prompt).text
-        if not full_response:
+        prompt = self._format_prompt(prompt)
+        inputs = self.model.tokenizer(prompt, return_tensors="pt").to(self.model.device)
+        
+        streamed_response = ""
+        for output in self.model.generate(**inputs, max_length=self.num_output, streaming=True):
+            new_token = self.model.tokenizer.decode(output[0], skip_special_tokens=True)
+            if new_token:
+                streamed_response += new_token
+                yield CompletionResponse(text=streamed_response, delta=new_token)
+        
+        if not streamed_response:
             yield CompletionResponse(text="No response generated.", delta="No response generated.")
-        else:
-            for token in full_response:
-                yield CompletionResponse(text=token, delta=token)
 
 # for transformers 1
 """class GemmaLLMInterface(CustomLLM):