Update app.py

app.py

@@ -95,7 +95,7 @@ class DocumentRetrievalAndGeneration:
         messages = [{"role": "user", "content": prompt}]
         encodeds = self.llm.tokenizer.apply_chat_template(messages, return_tensors="pt")
         model_inputs = encodeds.to(self.llm.device)
-
+        
         # Perform inference and measure time
         start_time = datetime.now()
         generated_ids = self.llm.model.generate(model_inputs, max_new_tokens=1000, do_sample=True)
@@ -104,11 +104,22 @@ class DocumentRetrievalAndGeneration:
         # Decode and return output
         decoded = self.llm.tokenizer.batch_decode(generated_ids)
         generated_response = decoded[0]
+        match1 = re.search(r'\[/INST\](.*?)</s>', generated_response, re.DOTALL)
+        match2 = re.search(r'Solution:(.*?)</s>', text, re.DOTALL | re.IGNORECASE)
+        if match1:
+            solution_text = match1.group(1).strip()
+            print(solution_text)
+        elif match2:
+            solution_text = match2.group(1).strip()
+            print(solution_text)
+            
+        else:
+            solution_text=generated_response
         print("Generated response:", generated_response)
         print("Time elapsed:", elapsed_time)
         print("Device in use:", self.llm.device)
 
-        return generated_response, content
+        return solution_text, content
 
     def qa_infer_gradio(self, query):
         response = self.query_and_generate_response(query)