Update app.py

app.py

@@ -1,66 +1,139 @@
 import gradio as gr
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import spaces 
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+import spaces
 
-# Load model & tokenizer
 MODEL_NAME = "ubiodee/Plutus_Tutor_new"
 
-# Initialize tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-model.eval()
-
-# Response generation function
-@spaces.GPU
-def generate_response(personality, level, topic):
-    # Construct a structured prompt incorporating user selections
-    full_prompt = (
+# --------- Lightweight utilities ----------
+def build_prompt(personality, level, topic):
+    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."
     )
-    
-    inputs = tokenizer(full_prompt, return_tensors="pt").to("cuda")
 
-    with torch.no_grad():
+def _ensure_tokenizer():
+    tok = AutoTokenizer.from_pretrained(MODEL_NAME)
+    if tok.pad_token_id is None:
+        tok.pad_token = tok.eos_token
+    return tok
+
+# 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=250,
+            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
+
+@spaces.GPU
+def generate_gpu(personality, level, topic, max_new_tokens=250):
+    """
+    Runs ONLY under a granted GPU.
+    Loads the model in 4-bit to fit ZeroGPU VRAM, generates, then frees VRAM.
+    """
+    tokenizer = _ensure_tokenizer()
+    prompt = build_prompt(personality, level, topic)
+
+    # Prefer 4-bit to minimize VRAM on ZeroGPU
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            MODEL_NAME,
+            load_in_4bit=True,
+            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.1,
             top_p=0.1,
             do_sample=True,
             eos_token_id=tokenizer.eos_token_id,
             pad_token_id=tokenizer.pad_token_id,
         )
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    # Remove the full prompt from the output to return only the answer
-    if response.startswith(full_prompt):
-        response = response[len(full_prompt):].strip()
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    if text.startswith(prompt):
+        text = text[len(prompt):].strip()
+
+    # Free VRAM ASAP
+    try:
+        del model
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    except Exception:
+        pass
+
+    return text
 
-    return response
+def orchestrator(personality, level, topic):
+    # Don’t run until all selections are made
+    if not personality or not level or not topic:
+        return "Select your personality, expertise, and topic to get a tailored explanation."
 
-# Gradio interface
+    # Try GPU path first; if ZeroGPU refuses/throws, fallback to CPU
+    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)}"
+    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)}"
+
+# --------- Gradio UI ----------
 with gr.Blocks(theme="default") as iface:
     gr.Markdown(
-        """
-        # Cardano Plutus AI Assistant
-        Select your learning personality, expertise level, and topic to get a tailored explanation about Plutus or Cardano.
-        The content will be generated automatically upon selection.
-        """
+        "## Cardano Plutus AI Assistant\n"
+        "Choose your **Learning Personality**, **Expertise Level**, and **Topic**. "
+        "An answer will be generated automatically."
     )
-    
+
     with gr.Row():
         personality = gr.Dropdown(
             choices=["Dyslexic", "Autistic", "Expressive"],
-            label="Select Your Learning Personality",
-            value="Expressive"
+            label="Learning Personality",
+            value=None,
+            allow_custom_value=False,
+            scale=1
         )
         level = gr.Dropdown(
             choices=["Beginner", "Intermediate", "Advanced"],
-            label="Select Your Expertise Level",
-            value="Beginner"
+            label="Expertise Level",
+            value=None,
+            allow_custom_value=False,
+            scale=1
         )
         topic = gr.Dropdown(
             choices=[
@@ -75,31 +148,31 @@ with gr.Blocks(theme="default") as iface:
                 "On-Chain Constraints",
                 "Plutus Core",
                 "Transaction Validation",
-                "Cardano Node Integration"
+                "Cardano Node Integration",
             ],
-            label="Select Topic",
-            value="Plutus Basics"
+            label="Topic",
+            value=None,
+            allow_custom_value=False,
+            scale=2
         )
-    
-    output = gr.Textbox(label="Model Response")
-
-    # Trigger generation on any dropdown change
-    personality.change(
-        fn=generate_response,
-        inputs=[personality, level, topic],
-        outputs=output
-    )
-    level.change(
-        fn=generate_response,
-        inputs=[personality, level, topic],
-        outputs=output
-    )
-    topic.change(
-        fn=generate_response,
-        inputs=[personality, level, topic],
-        outputs=output
-    )
 
-# Launch the app (Hugging Face Spaces handles this automatically)
+    with gr.Row():
+        regen = gr.Button("🔁 Regenerate")
+    output = gr.Textbox(label="Model Response", lines=12, interactive=False, show_copy_button=True)
+
+    # 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."
+
+    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)
+
+# Enable request queueing (helps with ZeroGPU scheduling)
+iface.queue(concurrency_count=1, max_size=8)
+
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    iface.launch(server_name="0.0.0.0", server_port=7860)