Update src/generation.py

src/generation.py

@@ -2,7 +2,11 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 class ResponseGenerator:
     def __init__(self, model_name="distilgpt2"):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, clean_up_tokenization_spaces=True)
+        # Set a distinct pad token
+        if self.tokenizer.pad_token is None:
+            self.tokenizer.pad_token = self.tokenizer.eos_token
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
         self.model = AutoModelForCausalLM.from_pretrained(model_name)
         self.prompt_template = """
 You are a customer support chatbot for Rupeia, a financial platform. Provide accurate, concise answers about Investments, Goals, Benefits, Gadgets, and News & Blogs. Use the context and history to respond naturally. If unsure, say: "I’m not sure about that. Could you clarify or ask about Rupeia features?"
@@ -15,7 +19,13 @@ Answer: """
     def generate_response(self, user_input, context, history):
         history_str = "\n".join([f"User: {h[0]}\nBot: {h[1]}" for h in history[-3:]])
         prompt = self.prompt_template.format(context=context, history=history_str, user_input=user_input)
-        inputs = self.tokenizer.encode(prompt, return_tensors="pt")
-        outputs = self.model.generate(inputs, max_length=200, pad_token_id=self.tokenizer.eos_token_id)
-        response = self.tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True).strip()
+        inputs = self.tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        # Ensure attention mask is passed
+        outputs = self.model.generate(
+            inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            max_length=200,
+            pad_token_id=self.tokenizer.pad_token_id
+        )
+        response = self.tokenizer.decode(outputs[:, inputs["input_ids"].shape[-1]:][0], skip_special_tokens=True).strip()
         return response if response else context or "I’m not sure about that. Could you clarify or ask about Rupeia features?"