Update app.py

app.py

@@ -58,7 +58,19 @@ def generate_response(user_query, relevant_segment):
 
         # Encode the input and generate a response
         input_ids = tokenizer.encode(user_message, return_tensors='pt')
-        output = model.generate(input_ids, max_length=150, num_return_sequences=1)
+        
+        # Create attention mask (1 for real tokens, 0 for padding tokens)
+        attention_mask = (input_ids != tokenizer.pad_token_id).long()
+
+        # Generate the response using the model
+        output = model.generate(
+            input_ids,
+            attention_mask=attention_mask,
+            max_length=150,
+            num_return_sequences=1,
+            pad_token_id=tokenizer.eos_token_id  # Set pad_token_id to eos_token_id
+        )
+
         output_text = tokenizer.decode(output[0], skip_special_tokens=True)
 
         # Append assistant's message to messages list for context