Update app.py

app.py

@@ -71,7 +71,7 @@ class DocumentRetrievalAndGeneration:
         return generate_text
     def initialize_llm2(self,model_id):
         try:
-            pipe = pipeline("text-generation", model="google/gemma-2-9b")
+            client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
         except:
             try:
                 pipe = pipeline("text-generation", model="microsoft/Phi-3-mini-4k-instruct", trust_remote_code=True)
@@ -123,7 +123,7 @@ class DocumentRetrievalAndGeneration:
                 print("############################")
             else:
                 print(f"Index {idx} is out of bounds. Skipping.")
-
+        # {query}
         prompt = f"""<s>
         You are a knowledgeable assistant with access to a comprehensive database. 
         I need you to answer my question and provide related information in a specific format.
@@ -132,7 +132,7 @@ class DocumentRetrievalAndGeneration:
         Include a final answer without additional comments, sign-offs, or extra phrases. Be direct and to the point.
         content
         Here's my question:
-        Query:{query}
+        Query:
         Solution==>
         RETURN ONLY SOLUTION . IF THEIR IS NO ANSWER RELATABLE IN RETRIEVED CHUNKS , RETURN " NO SOLUTION AVAILABLE"
         IF THE QUERY AND THE RETRIEVED CHUNKS DO NOT CORRELATE MEANINGFULLY, OR IF THE QUERY IS NOT RELEVANT TO TDA2 OR RELATED TOPICS, THEN "NO SOLUTION AVAILABLE."
@@ -151,11 +151,15 @@ class DocumentRetrievalAndGeneration:
         Solution:"NO SOLUTION AVAILABLE"
         </s>
         """
-        messages = [
-            {"role": "user", "content": prompt},
-        ]
-        
-        generated_response=pipe(messages)
+        messages = [{"role": "system", "content": prompt}]
+        messages.append({"role": "user", "content": message})
+        response = ""
+
+        for message in client.chat_completion(messages,max_tokens=2048,stream=True,temperature=0.7):
+            token = message.choices[0].delta.content
+            response += token
+        # yield response
+        generated_response=response
         # messages = [{"role": "user", "content": prompt}]
         # encodeds = self.llm.tokenizer.apply_chat_template(messages, return_tensors="pt")
         # model_inputs = encodeds.to(self.llm.device)