Update app.py

app.py

@@ -1,87 +1,66 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
-from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-MODEL_ID = "microsoft/phi-3-mini-4k-instruct"
+MODEL_ID = "microsoft/phi-2"  # lighter → IMPORTANT
 
-# Load model + tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float32,   # safer for CPU
-    device_map="auto"
+    device_map="cpu",
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=True
 )
 
-# Chat function with streaming
-def chat(message, history):
-    # Format conversation
-    messages = []
-
-    for user, bot in history:
-        messages.append({"role": "user", "content": user})
-        messages.append({"role": "assistant", "content": bot})
+torch.set_num_threads(2)
 
-    messages.append({"role": "user", "content": message})
+def generate_reply(message, history):
+    # Convert history to prompt
+    prompt = ""
+    for msg in history:
+        if msg["role"] == "user":
+            prompt += f"User: {msg['content']}\n"
+        else:
+            prompt += f"Assistant: {msg['content']}\n"
 
-    # Apply chat template
-    prompt = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    prompt += f"User: {message}\nAssistant:"
 
-    streamer = TextIteratorStreamer(
-        tokenizer,
-        skip_prompt=True,
-        skip_special_tokens=True
-    )
+    inputs = tokenizer(prompt, return_tensors="pt")
 
-    generation_kwargs = dict(
+    outputs = model.generate(
         **inputs,
-        streamer=streamer,
-        max_new_tokens=150,     # keep small for speed
+        max_new_tokens=80,
         temperature=0.7,
         do_sample=True
     )
 
-    thread = Thread(target=model.generate, kwargs=generation_kwargs)
-    thread.start()
+    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    reply = reply.split("Assistant:")[-1].strip()
 
-    partial_text = ""
-    for new_token in streamer:
-        partial_text += new_token
-        yield partial_text
+    return reply
 
 
-# Gradio UI
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("## ⚡ Phi-3 Mini Chatbot (Fast HF Space)")
+def chat(message, history):
+    if history is None:
+        history = []
 
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(placeholder="Type your message...")
-    clear = gr.Button("Clear")
+    history.append({"role": "user", "content": message})
 
-    def user_input(user_message, history):
-        return "", history + [[user_message, ""]]
+    reply = generate_reply(message, history)
 
-    def bot_response(history):
-        user_message = history[-1][0]
+    history.append({"role": "assistant", "content": reply})
 
-        bot_reply = ""
-        for chunk in chat(user_message, history[:-1]):
-            bot_reply = chunk
-            history[-1][1] = bot_reply
-            yield history
+    return history
 
-    msg.submit(user_input, [msg, chatbot], [msg, chatbot], queue=False).then(
-        bot_response, chatbot, chatbot
-    )
 
-    clear.click(lambda: None, None, chatbot, queue=False)
+with gr.Blocks() as demo:
+    gr.Markdown("## ⚡ Fast Phi-2 Chatbot (HF Free Tier)")
+
+    chatbot = gr.Chatbot(type="messages")
+    msg = gr.Textbox(placeholder="Type something...")
+    clear = gr.Button("Clear")
+
+    msg.submit(chat, [msg, chatbot], chatbot)
+    clear.click(lambda: [], None, chatbot)
 
-demo.queue()
 demo.launch()