Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 import torch
+import threading
 
 MODEL = "wmaousley/MiniCrit-1.5B"
 
@@ -8,24 +9,74 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
     torch_dtype=torch.float16,
-    device_map="cpu"   # Spaces CPU runtime
+    device_map="cpu"
 )
 
-def chat_fn(user_input):
-    inputs = tokenizer(user_input, return_tensors="pt")
-    outputs = model.generate(
+def generate_stream(prompt):
+    """Streaming generator."""
+    inputs = tokenizer(prompt, return_tensors="pt")
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+    
+    generation_kwargs = dict(
         **inputs,
         max_new_tokens=200,
         temperature=0.7,
         do_sample=True,
+        streamer=streamer
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-demo = gr.Interface(
-    fn=chat_fn,
-    inputs=gr.Textbox(label="Input"),
-    outputs=gr.Textbox(label="Response"),
-    title="MiniCrit-1.5B Chat"
-)
+    thread = threading.Thread(target=model.generate, kwargs=generation_kwargs)
+    thread.start()
+
+    for new_token in streamer:
+        yield new_token
+
+
+def chat_fn(message, history):
+    """Formats chat history + generates streaming reply."""
+    # Build conversation text
+    conversation = ""
+    for user, bot in history:
+        conversation += f"User: {user}\nMiniCrit: {bot}\n"
+    conversation += f"User: {message}\nMiniCrit:"
+
+    # Stream tokens
+    reply = ""
+    for token in generate_stream(conversation):
+        reply += token
+        yield reply
+
+
+# -------- UI --------
+
+with gr.Blocks(theme=gr.themes.Base()) as demo:
+
+    gr.Markdown(
+        """
+        <h1 style='text-align:center; color:#00eaff;'>
+            MiniCrit-1.5B Chat UI 🚀
+        </h1>
+        <p style='text-align:center; color:gray;'>Enhanced Streaming Interface</p>
+        """
+    )
+
+    chatbox = gr.Chatbot(
+        label="MiniCrit-1.5B",
+        height=500
+    )
+
+    with gr.Row():
+        msg = gr.Textbox(
+            placeholder="Ask something...",
+            label="Message",
+            scale=10
+        )
+        send = gr.Button("Send", variant="primary")
+        clear = gr.Button("Clear")
+
+    send.click(chat_fn, [msg, chatbox], chatbox)
+    send.click(lambda: "", None, msg)
+    clear.click(lambda: [], None, chatbox)
+
 
 demo.launch(debug=True)