Update rag.py

rag.py

@@ -115,7 +115,7 @@ class FinancialChatbot:
     
     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 without requiring context. 
+        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."
 
@@ -129,7 +129,9 @@ class FinancialChatbot:
         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)
-        return self.qwen_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        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):