Update app.py

app.py

@@ -2,11 +2,13 @@ import gradio as gr
 from transformers import pipeline
 import torch
 
-# Using Qwen 0.5B - The fastest model for Free CPUs
+# 1. SETUP: Qwen 0.5B - The "King" of Free CPU models
+# This model is small enough to be fast but smart enough to have a personality.
 model_id = "Qwen/Qwen2.5-0.5B-Instruct"
 
-print("Loading Mochi to CPU...")
-# Load the text-generation pipeline
+print("Loading Mochi to CPU... Stay chill.")
+
+# Using float32 because CPUs handle it more stably than float16
 pipe = pipeline(
     "text-generation", 
     model=model_id, 
@@ -16,50 +18,56 @@ pipe = pipeline(
 
 def chat_fn(message, history):
     """
-    message: The current user input
-    history: List of [user_message, assistant_message] pairs
+    Fixed chat function to prevent 'Role' errors on 2nd+ messages.
     """
-    
-    # 1. THE SOUL: Define Mochi's personality
-    character_card = """You are Mochi, a chill, slightly sarcastic, but deeply loyal best friend. 
-    - Talk like a real person in their 20s: use lowercase, occasional slang (idk, cool, lol), and short sentences.
-    - Be supportive but 'roast' the user occasionally if they say something silly.
-    - NEVER sound like a corporate AI. No 'I am here to assist you'.
-    - Keep answers short and snappy."""
+    # THE SOUL: Define Mochi's personality
+    character_card = (
+        "You are Mochi, a chill, slightly sarcastic, but deeply loyal best friend. "
+        "Talk like a real person in their 20s: use lowercase, occasional slang, and short sentences. "
+        "Be supportive but roast the user occasionally. Keep answers short and snappy."
+    )
 
-    # 2. THE MEMORY: Build the conversation history
+    # THE MEMORY: Properly format history for the model
+    # Format: [ {'role': 'system', 'content': '...'}, {'role': 'user', 'content': '...'}, {'role': 'assistant', 'content': '...'} ]
     messages = [{"role": "system", "content": character_card}]
     
+    # history comes from Gradio as a list of [user_message, assistant_message]
     for user_msg, assistant_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": assistant_msg})
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if assistant_msg:
+            messages.append({"role": "assistant", "content": assistant_msg})
     
-    # Add the current message
+    # Add the newest message from the user
     messages.append({"role": "user", "content": message})
     
-    # 3. GENERATION: Tweak temperature for personality
-    # max_new_tokens=150 keeps it fast on CPU
-    generation = pipe(
-        messages, 
-        max_new_tokens=150, 
-        do_sample=True, 
-        temperature=0.85, 
-        top_p=0.9
-    )
-    
-    # Extract the text result
-    return generation[0]['generated_text'][-1]['content']
+    # THE GENERATION: Optimized for speed and personality
+    try:
+        generation = pipe(
+            messages, 
+            max_new_tokens=128, 
+            do_sample=True, 
+            temperature=0.85, 
+            top_p=0.9,
+            truncation=True # Prevents RAM overflow as chat gets longer
+        )
+        # Safely extract the content
+        return generation[0]['generated_text'][-1]['content']
+    except Exception as e:
+        print(f"Error: {e}")
+        return "my bad, my brain just glitched. can you say that again? lol"
 
-# 4. THE UI: Custom "Buddy App" Theme
+# THE UI: Soft Pink Theme for a 'Buddy' feel
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="pink", secondary_hue="slate")) as demo:
     gr.Markdown("# 🐾 Mochi AI")
-    gr.Markdown("*Your chill, unlimited best friend. No limits, no filters.*")
+    gr.Markdown("*Your chill, unlimited best friend. No tokens, no limits.*")
     
+    # ChatInterface automatically handles the 'history' list for us
     chat_interface = gr.ChatInterface(
         fn=chat_fn,
-        examples=["hey mochi, what's up?", "tell me a secret", "roast my outfit lol"],
-        cache_examples=False,
+        examples=["yo mochi, what's up?", "roast my life real quick", "idk what to eat today"],
+        cache_examples=False
     )
 
-# 5. THE API: This enables the "sellable" API endpoint automatically
+# Launch the app (This also enables the API automatically)
 demo.launch()