Update app.py

app.py

@@ -27,28 +27,28 @@ class ChatbotUI:
         # Retrieve relevant context from Pinecone
         retrieved_docs = self.retriever.retrieve_context(input_text, top_k=1)
         
-        # Preset instruction for the legal assistant
         preset = "You are a HR legal assistant. You do not respond as 'User' or pretend to be 'User'. You only respond once as 'Assistant'. Avoid Yes or No."
-        
-        # Combine retrieved context with the user query to form the prompt
-        prompt = f"{preset}\nContext: {retrieved_docs}\nUser: {input_text}\nResponse:"
-        logging.info(f"Generated prompt for TinyLlama: {prompt}")
 
-        # Generate response using the model
-        inputs = self.trainer.tokenizer(prompt, return_tensors="pt").input_ids
-        output = self.trainer.model.generate(inputs, max_length=2000, do_sample=True, top_p=0.95, temperature=0.1, repetition_penalty=1.2)
-        response = self.trainer.tokenizer.decode(output[0], skip_special_tokens=True)
-        logging.info("Response generated successfully.")
-        
-        # Ensure response does not repeat the prompt
-        if response.startswith(prompt):
-            response = response[len(prompt):].strip()
+        latest_prompt = f"{preset}\n\n### Context: {retrieved_docs.strip()}\n\n### User: {prompt.strip()}\n\n### Response:"
 
-        # Append the new turn to the conversation history
-        history.append((input_text, response))
-        logging.info(f"Updated conversation history: {history}")
-        
-        return history, ""
+        #latest_prompt = f"{preset}\n\n### User: {prompt.strip()}\n\n### Response:"
+
+        inputs = tokenizer(latest_prompt, return_tensors="pt").to(model.device)
+
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=2000,  # Reduce token size to optimize speed
+            do_sample=True,
+            top_p=0.95,
+            temperature=0.1,
+            repetition_penalty=1.2,
+        )
+
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+        if response.startswith(latest_prompt):
+            response = response[len(latest_prompt):].strip()
+    
+            return response
 
     def clear_conversation(self):
         """Clears the entire conversation history and resets the input box."""