Update app.py

app.py

@@ -1,57 +1,75 @@
+import threading
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 import spaces
 
 MODEL_NAME = "ubiodee/Plutus_Tutor_new"
 
-# --------- Lightweight utilities ----------
+# ---------------- Utilities ----------------
+
+_TOKENIZER = None
+def get_tokenizer():
+    global _TOKENIZER
+    if _TOKENIZER is None:
+        tok = AutoTokenizer.from_pretrained(MODEL_NAME)
+        if tok.pad_token_id is None:
+            tok.pad_token = tok.eos_token
+        _TOKENIZER = tok
+    return _TOKENIZER
+
 def build_prompt(personality, level, topic):
+    # Conversational, clear, and bounded length to reduce truncation
     return (
-        f"You are a Plutus AI Assistant tailored for a {personality} learner "
-        f"at {level} level, focusing on {topic}. Provide a clear, concise, "
-        f"and tailored explanation of {topic}, suitable for the specified personality and expertise level."
+        f"You are a friendly Plutus AI tutor for a {personality} learner at {level} level.\n"
+        f"Topic: {topic}\n\n"
+        f"Explain in a conversational tone, with simple language and concrete examples.\n"
+        f"Keep it focused and complete in about 120–180 words.\n"
+        f"End with a single-sentence takeaway starting with 'Takeaway:'.\n"
     )
 
-def _ensure_tokenizer():
-    tok = AutoTokenizer.from_pretrained(MODEL_NAME)
-    if tok.pad_token_id is None:
-        tok.pad_token = tok.eos_token
-    return tok
+# ---------------- Streaming paths ----------------
+
+def from_cpu_stream(prompt, max_new_tokens=200):
+    tokenizer = get_tokenizer()
+    # Load on CPU (fallback)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+    model.eval()
 
-# CPU fallback (slow, but prevents total failure)
-def generate_cpu(personality, level, topic, max_new_tokens=250):
-    tokenizer = _ensure_tokenizer()
-    prompt = build_prompt(personality, level, topic)
     inputs = tokenizer(prompt, return_tensors="pt")
-    # Small settings for CPU to avoid long stalls
-    with torch.inference_mode():
-        model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)  # CPU load
-        model.eval()
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=min(max_new_tokens, 128),
-            temperature=0.2,
-            top_p=0.9,
-            do_sample=True,
-            eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.pad_token_id,
-        )
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    if text.startswith(prompt):
-        text = text[len(prompt):].strip()
-    return text
+
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+    gen_kwargs = dict(
+        **inputs,
+        max_new_tokens=max_new_tokens,
+        temperature=0.2,
+        top_p=0.9,
+        do_sample=True,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id,
+        streamer=streamer,
+    )
+
+    t = threading.Thread(target=model.generate, kwargs=gen_kwargs)
+    t.start()
+
+    buffer = ""
+    for piece in streamer:
+        buffer += piece
+        # Trim echoed prompt once
+        if buffer.startswith(prompt):
+            buffer = buffer[len(prompt):].lstrip()
+        yield buffer
 
 @spaces.GPU
-def generate_gpu(personality, level, topic, max_new_tokens=1000):
+def _gpu_generate_stream(prompt, max_new_tokens):
     """
-    Runs ONLY under a granted GPU.
-    Loads the model in 4-bit to fit ZeroGPU VRAM, generates, then frees VRAM.
+    This function runs ONLY when ZeroGPU grants a GPU.
+    We create a nested generator so Gradio can stream tokens.
     """
-    tokenizer = _ensure_tokenizer()
-    prompt = build_prompt(personality, level, topic)
+    tokenizer = get_tokenizer()
 
-    # Prefer 4-bit to minimize VRAM on ZeroGPU
+    # Prefer 4-bit to reduce VRAM; if not available, fall back to fp16
     try:
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
@@ -59,33 +77,38 @@ def generate_gpu(personality, level, topic, max_new_tokens=1000):
             device_map="auto",
         )
     except Exception:
-        # If 4-bit isn’t available for this arch, fallback to fp16 on GPU
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
             torch_dtype=torch.float16,
             device_map="auto",
         )
-
     model.eval()
 
     device = next(model.parameters()).device
     inputs = tokenizer(prompt, return_tensors="pt")
     inputs = {k: v.to(device) for k, v in inputs.items()}
 
-    with torch.inference_mode():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=max_new_tokens,
-            temperature=0.3,
-            top_p=0.3,
-            do_sample=True,
-            eos_token_id=tokenizer.eos_token_id,
-            pad_token_id=tokenizer.pad_token_id,
-        )
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+    gen_kwargs = dict(
+        **inputs,
+        max_new_tokens=max_new_tokens,
+        temperature=0.15,
+        top_p=0.9,
+        do_sample=True,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id,
+        streamer=streamer,
+    )
 
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    if text.startswith(prompt):
-        text = text[len(prompt):].strip()
+    t = threading.Thread(target=model.generate, kwargs=gen_kwargs)
+    t.start()
+
+    buffer = ""
+    for piece in streamer:
+        buffer += piece
+        if buffer.startswith(prompt):
+            buffer = buffer[len(prompt):].lstrip()
+        yield buffer
 
     # Free VRAM ASAP
     try:
@@ -95,24 +118,28 @@ def generate_gpu(personality, level, topic, max_new_tokens=1000):
     except Exception:
         pass
 
-    return text
-
-def orchestrator(personality, level, topic):
-    # Don’t run until all selections are made
+def orchestrator_stream(personality, level, topic):
     if not personality or not level or not topic:
-        return "Select your personality, expertise, and topic to get a tailored explanation."
+        yield "Select your personality, expertise, and topic to get a tailored explanation."
+        return
 
-    # Try GPU path first; if ZeroGPU refuses/throws, fallback to CPU
+    prompt = build_prompt(personality, level, topic)
+
+    # Try GPU streaming first
     try:
-        return generate_gpu(personality, level, topic)
-    except RuntimeError as e:
-        # Typical ZeroGPU worker errors show here – fall back gracefully
-        return f"(GPU unavailable, using CPU fallback)\n\n{generate_cpu(personality, level, topic)}"
+        for chunk in _gpu_generate_stream(prompt, max_new_tokens=240):
+            yield chunk
+        return
     except Exception as e:
-        # Any other unexpected issue – try CPU anyway
-        return f"(GPU error: {type(e).__name__})\n\n{generate_cpu(personality, level, topic)}"
+        # Log server-side; do NOT show to user
+        print(f"[GPU fallback] {type(e).__name__}: {e}")
+
+    # CPU streaming fallback
+    for chunk in from_cpu_stream(prompt, max_new_tokens=200):
+        yield chunk
+
+# ---------------- Gradio UI ----------------
 
-# --------- Gradio UI ----------
 with gr.Blocks(theme="default") as iface:
     gr.Markdown(
         "## Cardano Plutus AI Assistant\n"
@@ -126,14 +153,14 @@ with gr.Blocks(theme="default") as iface:
             label="Learning Personality",
             value=None,
             allow_custom_value=False,
-            scale=1
+            scale=1,
         )
         level = gr.Dropdown(
             choices=["Beginner", "Intermediate", "Advanced"],
             label="Expertise Level",
             value=None,
             allow_custom_value=False,
-            scale=1
+            scale=1,
         )
         topic = gr.Dropdown(
             choices=[
@@ -153,26 +180,29 @@ with gr.Blocks(theme="default") as iface:
             label="Topic",
             value=None,
             allow_custom_value=False,
-            scale=2
+            scale=2,
         )
 
     with gr.Row():
-        regen = gr.Button("🔁 Regenerate")
-    output = gr.Textbox(label="Model Response", lines=12, interactive=False, show_copy_button=True)
+        regen = gr.Button("🔁 Generate")
 
-    # Auto-generate when any dropdown changes (only once all three have values)
-    def _maybe_generate(p, l, t):
-        if p and l and t:
-            return orchestrator(p, l, t)
-        return "Select your personality, expertise, and topic to get a tailored explanation."
+    output = gr.Textbox(
+        label="Model Response",
+        lines=12,
+        interactive=False,
+        show_copy_button=True
+    )
 
-    personality.change(_maybe_generate, [personality, level, topic], output, queue=True)
-    level.change(_maybe_generate, [personality, level, topic], output, queue=True)
-    topic.change(_maybe_generate, [personality, level, topic], output, queue=True)
-    regen.click(orchestrator, [personality, level, topic], output, queue=True)
+    # Auto-generate whenever any dropdown changes (after all 3 are set)
+    def _maybe_stream(p, l, t):
+        return orchestrator_stream(p, l, t)
 
+    personality.change(_maybe_stream, [personality, level, topic], output, queue=True)
+    level.change(_maybe_stream, [personality, level, topic], output, queue=True)
+    topic.change(_maybe_stream, [personality, level, topic], output, queue=True)
+    regen.click(_maybe_stream, [personality, level, topic], output, queue=True)
 
-# Enable request queueing (simpler call for older Gradio)
+# Enable request queueing (use simple call for wider Gradio compatibility)
 iface.queue()
 
 if __name__ == "__main__":