Update app.py

app.py

@@ -1,52 +1,40 @@
-import os
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-
-# Use your Hugging Face token from the Secrets
-HF_TOKEN = os.environ.get("HF_TOKEN")
-
-# Model (lightweight, public, works on CPU)
-MODEL_NAME = "TheBloke/Gemma-3-270M-GGML"
-
-# Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=HF_TOKEN)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, use_auth_token=HF_TOKEN).to("cpu")
-
-history = []
-
-def chat_with_airi(user_msg):
-    global history
-    # build simple prompt
-    prompt = ""
-    for u, a in history[-5:]:
-        prompt += f"User: {u}\nAiri: {a}\n"
-    prompt += f"User: {user_msg}\nAiri:"
-
-    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
-    with torch.no_grad():
-        output = model.generate(
-            **inputs,
-            max_new_tokens=60,
-            do_sample=True,
-            top_p=0.9,
-            temperature=0.8,
-            pad_token_id=tokenizer.eos_token_id
-        )
-
-    reply = tokenizer.decode(output[0], skip_special_tokens=True)
-    reply = reply.split("Airi:", 1)[-1].strip()
-
-    history.append([user_msg, reply])
-    return history, ""
-
-# Gradio interface
-with gr.Blocks() as demo:
-    gr.HTML("<h2 style='text-align:center'>Airi — Mini Chat AI</h2>")
-    gr.HTML("<p style='text-align:center;color:#666;'>Small, Fast & Public Model</p>")
-
-    chat = gr.Chatbot()
-    msg = gr.Textbox(label="Talk to Airi…", placeholder="Write here…")
-    msg.submit(chat_with_airi, msg, [chat, msg])
-
-demo.launch()
+from transformers import pipeline
+
+# Initialize the chatbot model (publicly available, no auth needed)
+# Using a lightweight model for fast responses
+model_id = "facebook/blenderbot-400M-distill"
+chatbot = pipeline("conversational", model=model_id)
+
+def chat_function(message, history):
+    """
+    Simple chat function that takes user input and returns AI response
+    """
+    # Convert history to format expected by the model
+    conversation = []
+    for human, ai in history:
+        conversation.append({"role": "user", "content": human})
+        conversation.append({"role": "assistant", "content": ai})
+    
+    conversation.append({"role": "user", "content": message})
+    
+    # Generate response
+    response = chatbot(conversation)
+    return response[-1]["content"]
+
+# Create the chat interface
+demo = gr.ChatInterface(
+    fn=chat_function,
+    title="airi",
+    description="A friendly AI assistant ready to chat with you!",
+    theme=gr.themes.Soft(),
+    examples=[
+        "Hello! How are you?",
+        "Tell me something interesting",
+        "What's your favorite topic?"
+    ],
+    cache_examples=False,
+)
+
+if __name__ == "__main__":
+    demo.launch()