Update app.py

app.py

@@ -21,9 +21,20 @@ def chat(user_input):
         prompt += f"{speaker}: {msg}\n"
     prompt += f"User: {user_input}\nBot:"
 
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+    # Tokenize input with attention mask
+    inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs.input_ids
+    attention_mask = inputs.attention_mask
+
     with torch.no_grad():
-        output = model.generate(input_ids, max_new_tokens=100, do_sample=True, temperature=0.7)
+        output = model.generate(
+            input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=100,
+            do_sample=True,
+            temperature=0.7,
+            pad_token_id=tokenizer.eos_token_id
+        )
 
     full_reply = tokenizer.decode(output[0], skip_special_tokens=True)
     reply = full_reply.split("Bot:")[-1].strip()
@@ -48,6 +59,7 @@ with gr.Blocks() as demo:
     gr.Markdown("""# 🤖 Chat with Phi-2
 Small CPU-friendly chatbot using Hugging Face Transformers
 """)
+
     chatbot = gr.Markdown()
     with gr.Row():
         user_input = gr.Textbox(placeholder="Type your message here...", show_label=False)
@@ -62,6 +74,7 @@ Small CPU-friendly chatbot using Hugging Face Transformers
     clear_btn.click(clear, outputs=chatbot)
     export_btn.click(export_chat, outputs=file_out)
 
+    # Initialize chatbot display
     chatbot.value = format_chat(chat_history)
 
-demo.launch()
+demo.launch(share=True)