Update app.py

app.py

@@ -2,31 +2,58 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# Load tokenizer and model
+# Load model and tokenizer
 model_name = "microsoft/phi-3-mini-4k-instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
+model.to("cpu")  # CPU for Hugging Face Spaces
 
-# Move to CPU (safe for Hugging Face Spaces)
-model.to("cpu")
+# Optional: Custom system prompt (context for assistant behavior)
+SYSTEM_PROMPT = "You are a friendly and knowledgeable assistant who answers clearly and helpfully."
 
-# Basic chat function
-def generate_response(prompt):
+# Chat formatting using Phi-3's expected format
+def format_chat(system_prompt, history, user_input):
+    chat = ""
+    if system_prompt:
+        chat += f"<|system|>\n{system_prompt}\n"
+    for user, assistant in history:
+        chat += f"<|user|>\n{user}\n<|assistant|>\n{assistant}\n"
+    chat += f"<|user|>\n{user_input}\n<|assistant|>\n"
+    return chat
+
+# Inference function
+def chat_fn(message, chat_history):
+    prompt = format_chat(SYSTEM_PROMPT, chat_history, message)
     inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
     outputs = model.generate(
         **inputs,
-        max_new_tokens=256,
+        max_new_tokens=512,
         do_sample=True,
         temperature=0.7,
-        top_p=0.9
+        top_p=0.9,
+        pad_token_id=tokenizer.eos_token_id
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# Gradio UI
-gr.Interface(
-    fn=generate_response,
-    inputs=gr.Textbox(lines=5, label="Enter your question or prompt"),
-    outputs="text",
-    title="Phi-3 Mini Chat",
-    description="Powered by microsoft/phi-3-mini-4k-instruct"
-).launch()
+    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    # Extract only the latest assistant reply (after the last <|assistant|>)
+    if "<|assistant|>" in output_text:
+        response = output_text.split("<|assistant|>")[-1].strip()
+    else:
+        response = output_text.strip()
+
+    chat_history.append((message, response))
+    return "", chat_history
+
+# UI setup
+with gr.Blocks(theme=gr.themes.Base()) as demo:
+    gr.Markdown("## 🤖 Phi-3 Mini Chatbot\nTalk to a compact AI assistant powered by Microsoft's Phi-3 model.")
+    chatbot = gr.Chatbot(show_label=False)
+    msg = gr.Textbox(label="Type your message here...", placeholder="Ask me anything!", lines=2)
+    clear = gr.Button("Clear Chat")
+
+    state = gr.State([])  # Keeps chat history
+
+    msg.submit(chat_fn, [msg, state], [msg, chatbot])
+    clear.click(lambda: ([], []), None, [state, chatbot])
+
+demo.launch()