Upload folder using huggingface_hub

app.py

@@ -7,8 +7,7 @@ MODEL_ID = "Splashdude/smollm-chatbot"
 SYSTEM_PROMPT = (
     "You are a helpful, friendly AI assistant. "
     "You give clear, accurate, and conversational answers. "
-    "Keep responses concise unless the user asks for detail. "
-    "Remember what the user tells you and refer back to it when asked."
+    "Remember what the user tells you in this conversation."
 )
 
 model = None
@@ -27,58 +26,31 @@ def load_model():
     print("Model loaded!")
 
 
-def extract_text(content):
-    """Extract plain text from Gradio content (handles str, list, dict)."""
-    if isinstance(content, str):
-        return content
-    if isinstance(content, list):
-        parts = []
-        for item in content:
-            if isinstance(item, str):
-                parts.append(item)
-            elif isinstance(item, dict):
-                parts.append(item.get("text", item.get("content", str(item))))
-            else:
-                parts.append(str(item))
-        return " ".join(parts)
-    if isinstance(content, dict):
-        return content.get("text", content.get("content", str(content)))
-    return str(content)
-
-
-def build_messages(history, new_message):
-    """Convert Gradio history to chat messages list."""
-    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
-    for msg in history:
-        role = msg.get("role", "")
-        content = extract_text(msg.get("content", ""))
-        if role in ("user", "assistant") and content:
-            messages.append({"role": role, "content": content})
-    messages.append({"role": "user", "content": new_message})
-    return messages
-
-
-def generate_response(message, history):
+def generate_response(message, chat_history):
     if model is None:
         try:
             load_model()
         except Exception as e:
-            yield f"Error loading model: {e}"
+            chat_history.append({"role": "user", "content": message})
+            chat_history.append({"role": "assistant", "content": f"Error: {e}"})
+            yield chat_history, ""
             return
 
     if not message or not message.strip():
-        yield "Please type a message."
+        yield chat_history, ""
         return
 
-    try:
-        messages = build_messages(history, message)
-        text = tokenizer.apply_chat_template(
-            messages, tokenize=False, add_generation_prompt=True
-        )
-        inputs = tokenizer(text, return_tensors="pt")
-    except Exception as e:
-        yield f"Error building prompt: {e}"
-        return
+    chat_history.append({"role": "user", "content": message})
+    chat_history.append({"role": "assistant", "content": ""})
+
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    for msg in chat_history[:-1]:
+        messages.append({"role": msg["role"], "content": msg["content"]})
+
+    text = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=True
+    )
+    inputs = tokenizer(text, return_tensors="pt")
 
     streamer = TextIteratorStreamer(
         tokenizer, skip_prompt=True, skip_special_tokens=True
@@ -101,26 +73,67 @@ def generate_response(message, history):
     partial = ""
     for token in streamer:
         partial += token
-        yield partial
+        chat_history[-1]["content"] = partial
+        yield chat_history, ""
 
     thread.join()
 
 
-demo = gr.ChatInterface(
-    fn=generate_response,
-    title="AI Chatbot",
-    description="A fast conversational AI chatbot powered by SmolLM2-360M.",
-    examples=[
-        "Hello! How are you?",
-        "Tell me a joke.",
-        "What is the capital of France?",
-        "Explain gravity in simple terms.",
-        "Write a short poem about the ocean.",
-        "What are some tips for staying productive?",
-    ],
-    save_history=True,
-    editable=True,
-)
+def clear_chat():
+    return [], ""
+
+
+with gr.Blocks(title="AI Chatbot", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# AI Chatbot\nFast conversational AI powered by SmolLM2-360M.")
+
+    chatbot = gr.Chatbot(type="messages", height=500, show_copy_button=True, label="Chat")
+    chat_state = gr.State([])
+
+    with gr.Row():
+        msg = gr.Textbox(
+            placeholder="Type your message...",
+            show_label=False,
+            container=False,
+            scale=8,
+        )
+        submit = gr.Button("Send", variant="primary", scale=1)
+        clear = gr.Button("New Chat", scale=1)
+
+    gr.Examples(
+        examples=[
+            "Hello! How are you?",
+            "Tell me a joke.",
+            "What is the capital of France?",
+            "Explain gravity in simple terms.",
+        ],
+        inputs=msg,
+        label="Examples",
+    )
+
+    def user_submit(message, history):
+        for updated_history, _ in generate_response(message, history):
+            yield updated_history, "", updated_history
+
+    def bot_response(message, history):
+        for updated_history, _ in generate_response(message, history):
+            yield updated_history, updated_history
+
+    msg.submit(
+        user_submit,
+        [msg, chat_state],
+        [chatbot, msg, chat_state],
+        queue=True,
+    )
+
+    submit.click(
+        user_submit,
+        [msg, chat_state],
+        [chatbot, msg, chat_state],
+        queue=True,
+    )
+
+    clear.click(clear_chat, None, [chatbot, chat_state])
 
 if __name__ == "__main__":
+    demo.queue()
     demo.launch()