Update main.py

main.py

@@ -54,8 +54,10 @@ async def generate_response_chunks(prompt: str):
     else:
         input_ids = new_input_ids
 
+    attention_mask = input_ids.ne(tokenizer.pad_token_id).long()
     output_ids = model.generate(
         input_ids,
+        attention_mask=attention_mask,
         max_new_tokens=200,
         do_sample=True,
         top_p=0.9,
@@ -63,6 +65,7 @@ async def generate_response_chunks(prompt: str):
         pad_token_id=tokenizer.eos_token_id
     )
 
+
     chat_history_ids = output_ids  # update history
 
     response = tokenizer.decode(output_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)