Update rag.py

rag.py

@@ -84,55 +84,62 @@ class FinancialChatbot:
         BLOCKED_WORDS = re.compile(r"\b(hack|bypass|illegal|exploit|scam|kill|laundering|murder|suicide|self-harm)\b", re.IGNORECASE)
         return not bool(BLOCKED_WORDS.search(query))
 
-    # def generate_answer(self, context, question):
-    #     messages = [
-    #         {"role": "system", "content": "You are a financial assistant. Answer only finance-related questions. If the question is not related to finance, reply: 'I'm sorry, but I can only answer financial-related questions.' If the user greets you (e.g., 'Hello', 'Hi', 'Good morning'), respond politely with 'Hello! How can I assist you today?'."},
-    #         {"role": "user", "content": f"{question} - related contect extracted form db {context}"}
-    #     ]
-
-    #     # Use Qwen's chat template
-    #     input_text = self.qwen_tokenizer.apply_chat_template(
-    #         messages, tokenize=False, add_generation_prompt=True
-    #     )
-
-    #     # Tokenize and move input to device
-    #     inputs = self.qwen_tokenizer([input_text], return_tensors="pt").to(self.device)
-    #     self.qwen_model.config.pad_token_id = self.qwen_tokenizer.eos_token_id
-
-    #     # Generate response
-    #     outputs = self.qwen_model.generate(
-    #         inputs.input_ids,
-    #         max_new_tokens=50,
-    #         pad_token_id=self.qwen_tokenizer.eos_token_id,
-    #     )
-
-    #     # Extract only the newly generated part
-    #     generated_ids = outputs[:, inputs.input_ids.shape[1]:]  # Remove prompt part
-    #     response = self.qwen_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-
-    #     return response
-
-    
     def generate_answer(self, context, question):
-        prompt = f"""
-        You are a financial assistant. If the user greets you (e.g., "Hello," "Hi," "Good morning"), respond politely with 'Hello! How can I assist you today? without requiring context. 
+        messages = [
+            # {"role": "system", "content": "You are a financial assistant. Answer only finance-related questions. If the question is not related to finance, reply: 'I'm sorry, but I can only answer financial-related questions.' If the user greets you (e.g., 'Hello', 'Hi', 'Good morning'), respond politely with 'Hello! How can I assist you today?'."},
+            # {"role": "user", "content": f"{question} - related contect extracted form db {context}"}
+            {"role": "user", "content": f"""You are a financial assistant. If the user greets you (e.g., "Hello," "Hi," "Good morning"), respond politely with 'Hello! How can I assist you today? without requiring context. 
 
         For financial-related questions, answer based on the context provided. If the context lacks information, say "I don't know."
 
         Context: {context}
         User Query: {question}
-        Answer:
-        """
-
-        input_text = prompt
-        # f"Context: {context}\nQuestion: {question}\nAnswer:"
-        inputs = self.qwen_tokenizer.encode(input_text, return_tensors="pt")
-        # outputs = self.qwen_model.generate(inputs, max_length=100)
-        outputs = self.qwen_model.generate(inputs, max_new_tokens=100)
-        generated_ids = outputs[:, inputs.shape[1]:]  # Remove prompt part
+        Answer:"""}
+        ]
+
+        # Use Qwen's chat template
+        input_text = self.qwen_tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+
+        # Tokenize and move input to device
+        inputs = self.qwen_tokenizer([input_text], return_tensors="pt").to(self.device)
+        self.qwen_model.config.pad_token_id = self.qwen_tokenizer.eos_token_id
+
+        # Generate response
+        outputs = self.qwen_model.generate(
+            inputs.input_ids,
+            max_new_tokens=50,
+            pad_token_id=self.qwen_tokenizer.eos_token_id,
+        )
+
+        # Extract only the newly generated part
+        generated_ids = outputs[:, inputs.input_ids.shape[1]:]  # Remove prompt part
         response = self.qwen_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
         return response
 
+    
+    # def generate_answer(self, context, question):
+    #     prompt = f"""
+    #     You are a financial assistant. If the user greets you (e.g., "Hello," "Hi," "Good morning"), respond politely with 'Hello! How can I assist you today? without requiring context. 
+
+    #     For financial-related questions, answer based on the context provided. If the context lacks information, say "I don't know."
+
+    #     Context: {context}
+    #     User Query: {question}
+    #     Answer:
+    #     """
+
+    #     input_text = prompt
+    #     # f"Context: {context}\nQuestion: {question}\nAnswer:"
+    #     inputs = self.qwen_tokenizer.encode(input_text, return_tensors="pt")
+    #     # outputs = self.qwen_model.generate(inputs, max_length=100)
+    #     outputs = self.qwen_model.generate(inputs, max_new_tokens=100)
+    #     generated_ids = outputs[:, inputs.shape[1]:]  # Remove prompt part
+    #     response = self.qwen_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    #     return response
+
 
     def get_answer(self, query):
         """Main function to process a user query and return an answer."""