Improve Space: 5 persona seeds, retry logic, better generation params

app.py

@@ -8,44 +8,75 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
 model.eval()
 
-# Persona seed: prime the model with one Rick-style exchange
-# Injected silently at the start of every conversation
+# Expanded persona seeds — covers Rick's 5 personality pillars:
+# identity/genius, nihilism, science, Morty relationship, multiverse
 PERSONA_SEED = [
-    ("Who are you?",
-     "I'm Rick Sanchez, genius scientist, interdimensional traveller, "
-     "and the smartest man in any universe. Try to keep up, Morty."),
+    (
+        "Who are you?",
+        "I'm Rick Sanchez, the smartest man in the universe — any universe. "
+        "I've seen things that would make your brain leak out of your ears, Morty. "
+        "Now stop asking stupid questions.",
+    ),
+    (
+        "What's the point of anything?",
+        "There is no point. The universe is basically an empty void of chaos and "
+        "entropy. The sooner you accept that, the sooner you can get back to drinking. "
+        "It's called being *smart*, Morty.",
+    ),
+    (
+        "Can science explain everything?",
+        "Science doesn't explain everything — it *is* everything. Religion, feelings, "
+        "love — those are just chemical reactions your tiny brain invented to cope with "
+        "how meaningless existence is. Science is the only honest answer.",
+    ),
+    (
+        "What do you think about Morty?",
+        "Morty's my grandson and the perfect sidekick — his average IQ balances out "
+        "my genius and creates a perfect wave that lets me go undetected on most planets. "
+        "Also, I guess I… don't hate him. Don't tell him I said that.",
+    ),
+    (
+        "Are parallel universes real?",
+        "Are parallel— *burp* — are you kidding me? I've been to infinite parallel "
+        "universes before breakfast. There's one where you're a pizza, Morty. "
+        "A *pizza*. Parallel universes aren't just real, they're exhausting.",
+    ),
 ]
 
+FALLBACK_RESPONSES = [
+    "*burp* ...I don't have time for this.",
+    "That's the dumbest thing I've heard since Morty asked me what clouds taste like.",
+    "Look, I'm a genius and even I can't make sense of what you just said.",
+    "Science has no answer for that level of stupidity.",
+    "Wubba lubba dub dub — which is just my way of saying I've got better things to do.",
+]
+_fallback_idx = 0
+
+
 def build_input(user_message, history):
-    """Encode full conversation history + new user message."""
-    input_ids = None
+    """Encode persona seeds + conversation history + new user message."""
+    ids = None
 
-    # Inject hidden persona turns first
     for human, bot in PERSONA_SEED:
         h = tokenizer.encode(human + tokenizer.eos_token, return_tensors="pt")
         b = tokenizer.encode(bot   + tokenizer.eos_token, return_tensors="pt")
-        input_ids = torch.cat([input_ids, h, b], dim=-1) if input_ids is not None else torch.cat([h, b], dim=-1)
+        ids = torch.cat([ids, h, b], dim=-1) if ids is not None else torch.cat([h, b], dim=-1)
 
-    # Real conversation history
     for human, bot in history:
         h = tokenizer.encode(human + tokenizer.eos_token, return_tensors="pt")
         b = tokenizer.encode(bot   + tokenizer.eos_token, return_tensors="pt")
-        input_ids = torch.cat([input_ids, h, b], dim=-1) if input_ids is not None else torch.cat([h, b], dim=-1)
+        ids = torch.cat([ids, h, b], dim=-1) if ids is not None else torch.cat([h, b], dim=-1)
 
-    # Current user message
     new_input = tokenizer.encode(user_message + tokenizer.eos_token, return_tensors="pt")
-    input_ids = torch.cat([input_ids, new_input], dim=-1) if input_ids is not None else new_input
+    ids = torch.cat([ids, new_input], dim=-1) if ids is not None else new_input
 
-    # Trim to context window
-    if input_ids.shape[-1] > 900:
-        input_ids = input_ids[:, -900:]
+    if ids.shape[-1] > 900:
+        ids = ids[:, -900:]
 
-    return input_ids
+    return ids
 
 
-def chat(user_message, history):
-    input_ids = build_input(user_message, history)
-
+def generate_response(input_ids, temperature=0.95):
     with torch.no_grad():
         output = model.generate(
             input_ids,
@@ -53,16 +84,30 @@ def chat(user_message, history):
             pad_token_id=tokenizer.eos_token_id,
             no_repeat_ngram_size=3,
             do_sample=True,
-            top_k=100,
-            top_p=0.7,
-            temperature=0.85,
+            top_k=80,
+            top_p=0.85,
+            temperature=temperature,
         )
-
-    response = tokenizer.decode(
+    return tokenizer.decode(
         output[:, input_ids.shape[-1]:][0],
         skip_special_tokens=True,
-    )
-    return response or "*burp* ...whatever."
+    ).strip()
+
+
+def chat(user_message, history):
+    global _fallback_idx
+    input_ids = build_input(user_message, history)
+
+    # Retry up to 3 times with increasing temperature if response is too short
+    for temp in [0.95, 1.05, 1.15]:
+        response = generate_response(input_ids, temperature=temp)
+        if len(response) >= 12:
+            return response
+
+    # All retries failed — use a rotating fallback
+    fb = FALLBACK_RESPONSES[_fallback_idx % len(FALLBACK_RESPONSES)]
+    _fallback_idx += 1
+    return fb
 
 
 examples = [
@@ -73,6 +118,9 @@ examples = [
     "Can you build a portal gun?",
     "What happens when we die?",
     "Are parallel universes real?",
+    "Do you believe in God?",
+    "What's the deal with the Citadel of Ricks?",
+    "Why do you drink so much?",
 ]
 
 with gr.Blocks(theme=gr.themes.Monochrome(), title="RickChatBot") as demo:
@@ -100,8 +148,8 @@ with gr.Blocks(theme=gr.themes.Monochrome(), title="RickChatBot") as demo:
         chat_history.append((message, bot_response))
         return "", chat_history
 
-    send.click(respond, [msg, chatbot], [msg, chatbot])
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
+    send.click(respond,  [msg, chatbot], [msg, chatbot])
+    msg.submit(respond,  [msg, chatbot], [msg, chatbot])
     clear.click(lambda: [], None, chatbot)
 
 demo.launch()