Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-MODEL_NAME = "microsoft/DialoGPT-small"  # small, CPU-friendly
+MODEL_NAME = "microsoft/DialoGPT-small"  # small, CPU-friendly model [web:103]
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
@@ -41,15 +41,30 @@ Rules:
 def respond(message: str, history: list[dict]) -> str:
     """
     Gradio ChatInterface(type='messages') calls this as (message, history)
-    and expects JUST a reply string or message dict.[web:120]
+    where history is a list of dicts like:
+      {"role": "user" | "assistant", "content": "..."} [web:120]
+
+    We return a single reply string.
     """
 
-    # Build a simple one-turn prompt: system + latest user message.
-    prompt = (
-        f"System: {SHINCHAN_SYSTEM_PROMPT}\n\n"
-        f"User: {message}\n"
-        f"Shinchan:"
-    )
+    # Build a simple prompt: system + last few turns + latest message
+    lines = [f"System: {SHINCHAN_SYSTEM_PROMPT}"]
+
+    # keep prompt short: last 4 user/assistant messages max
+    trimmed = history[-8:] if history else []
+    for turn in trimmed:
+        role = turn.get("role")
+        content = turn.get("content", "")
+        if not content:
+            continue
+        if role == "user":
+            lines.append(f"User: {content}")
+        elif role == "assistant":
+            lines.append(f"Shinchan: {content}")
+
+    lines.append(f"User: {message}")
+    lines.append("Shinchan:")
+    prompt = "\n".join(lines)
 
     inputs = tokenizer(prompt, return_tensors="pt")
 
@@ -63,29 +78,28 @@ def respond(message: str, history: list[dict]) -> str:
             temperature=0.9,
         )
 
-    full_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-
-    # Extract text after the last "Shinchan:" marker.
-    reply = ""
-    if "Shinchan:" in full_text:
-        reply = full_text.split("Shinchan:")[-1].strip()
-    else:
-        # Fallback: everything after the original prompt length
-        reply = full_text[len(prompt):].strip()
+    # IMPORTANT: only take the newly generated tokens AFTER the prompt
+    input_len = inputs["input_ids"].shape[1]
+    generated_ids = output_ids[0][input_len:]
+    reply = tokenizer.decode(generated_ids, skip_special_tokens=True).strip()
 
-    # If model somehow returns empty or just punctuation, use a safe fallback.
-    if not reply or len(reply.replace(".", "").replace("!", "").strip()) < 3:
+    # If it's somehow empty, use a fallback so she never sees blank
+    if not reply:
         reply = (
             "Heeey, it’s Shinchan! 😂 I heard you, even if my brain glitched for a second. "
             "Tell me more, I’m listening. 🌻"
         )
 
+    # Optional: hard limit on length
+    if len(reply) > 400:
+        reply = reply[:380].rstrip() + "…"
+
     return reply
 
 
 demo = gr.ChatInterface(
     fn=respond,
-    type="messages",  # use new messages format internally[web:120]
+    type="messages",  # uses role/content dicts internally [web:120]
     title="Shinchan for Ruru",
     description="Private Shinchan-style chat for Ruru.",
 )