added new class

backend.py

@@ -33,9 +33,10 @@ model.eval()
 #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, tokenizer=tokenizer)"""
 Settings.embed_model = InstructorEmbedding(model_name="hkunlp/instructor-base")
-Settings.llm  = GemmaLLMInterface(model=model, tokenizer=tokenizer)
-#Settings.llm  = llm
+Settings.llm = GemmaLLMInterface(model_id="google/gemma-2-2b-it")
 
 
 ############################---------------------------------
@@ -57,7 +58,8 @@ def build_index():
 
 
 @spaces.GPU(duration=20)
-async def handle_query(query_str, chathistory):
+def handle_query(query_str, chathistory):
+    
     index = build_index()
   
     qa_prompt_str = (
@@ -73,32 +75,45 @@ async def handle_query(query_str, chathistory):
     chat_text_qa_msgs = [
         (
             "system",
-            "Sei un assistente italiano di nome Ossy che risponde solo alle domande o richieste pertinenti.",
+            "Sei un assistente italiano di nome Ossy che risponde solo alle domande o richieste pertinenti. ",
         ),
         ("user", qa_prompt_str),
     ]
     text_qa_template = ChatPromptTemplate.from_messages(chat_text_qa_msgs)
 
     try:
+        # Create a streaming query engine
+        """query_engine = index.as_query_engine(text_qa_template=text_qa_template, streaming=False, similarity_top_k=1)
+        
+        # Execute the query
+        streaming_response = query_engine.query(query_str)
+        
+        r = streaming_response.response
+        cleaned_result = r.replace("<end_of_turn>", "").strip()
+        yield cleaned_result"""
+      
+        # Stream the response
+        """outputs = []
+        for text in streaming_response.response_gen:
+            
+            outputs.append(str(text))
+            yield "".join(outputs)"""
+          
         memory = ChatMemoryBuffer.from_defaults(token_limit=1500)
         chat_engine = index.as_chat_engine(
-            chat_mode="context",
-            memory=memory,
-            system_prompt=(
-                "Sei un assistente italiano di nome Ossy che risponde solo alle domande o richieste pertinenti."
-            ),
+        chat_mode="context",
+        memory=memory,
+        system_prompt=(
+            "Sei un assistente italiano di nome Ossy che risponde solo alle domande o richieste pertinenti. "
+          ),
         )
         
-        # Stream the response
         response = chat_engine.stream_chat(query_str)
-        outputs = []
-        
-        async for token in response.response_gen:
-            outputs.append(token)
-            yield "".join(outputs)
+        #response = chat_engine.chat(query_str)
+        for token in response.response_gen:
+            yield token
+          
         
-    except StopAsyncIteration:
-        yield "No more responses to stream."
     except Exception as e:
         yield f"Error processing query: {str(e)}"
 

interface.py

@@ -6,8 +6,71 @@ import torch
 from transformers import TextIteratorStreamer
 from threading import Thread
 
-
+# for transformers 2
 class GemmaLLMInterface(CustomLLM):
+    def __init__(self, model_id: str = "google/gemma-2-2b-it", context_window: int = 8192, num_output: int = 2048):
+        self.model_id = model_id
+        self.context_window = context_window
+        self.num_output = num_output
+        
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_id,
+            device_map="auto",
+            torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+        )
+        self.model.eval()
+
+    def _format_prompt(self, message: str) -> str:
+        return f"<start_of_turn>user\n{message}<end_of_turn>\n<start_of_turn>model\n"
+
+    @property
+    def metadata(self) -> LLMMetadata:
+        return LLMMetadata(
+            context_window=self.context_window,
+            num_output=self.num_output,
+            model_name=self.model_id,
+        )
+
+    @llm_completion_callback()
+    def complete(self, prompt: str, **kwargs: Any) -> CompletionResponse:
+        formatted_prompt = self._format_prompt(prompt)
+        inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.model.device)
+        
+        with torch.no_grad():
+            outputs = self.model.generate(
+                **inputs,
+                max_new_tokens=self.num_output,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95,
+            )
+        
+        response = self.tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
+        return CompletionResponse(text=response)
+
+    @llm_completion_callback()
+    def stream_complete(self, prompt: str, **kwargs: Any) -> CompletionResponseGen:
+        formatted_prompt = self._format_prompt(prompt)
+        inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.model.device)
+        
+        response = ""
+        with torch.no_grad():
+            for output in self.model.generate(
+                **inputs,
+                max_new_tokens=self.num_output,
+                do_sample=True,
+                temperature=0.7,
+                top_p=0.95,
+                streamer=True,
+            ):
+                token = self.tokenizer.decode(output, skip_special_tokens=True)
+                response += token
+                yield CompletionResponse(text=response, delta=token)
+
+
+# for transformers 1
+"""class GemmaLLMInterface(CustomLLM):
     model: Any
     tokenizer: Any
     context_window: int = 8192
@@ -76,4 +139,4 @@ class GemmaLLMInterface(CustomLLM):
             for new_token in streamer:
                 yield CompletionResponse(text=new_token)
         except StopIteration:
-            return
+            return"""

requirements.txt

@@ -6,7 +6,7 @@ llama-index-embeddings-instructor
 sentence-transformers==2.2.2
 llama-index-readers-web
 llama-index-readers-file
-gradio==4.17.0
+gradio
 transformers
 llama-cpp-agent>=0.2.25
 setuptools