Update app.py

app.py

@@ -1,204 +1,46 @@
 import gradio as gr
 import torch
-import time
 from transformers import AutoModelForCausalLM, AutoTokenizer
-import logging
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
 
+# Load model & tokenizer
 MODEL_NAME = "ubiodee/Plutus_Tutor_new"
 
-# ------------ Model and Tokenizer cache ------------
-_TOKENIZER = None
-_MODEL = None
-
-def get_tokenizer():
-    global _TOKENIZER
-    if _TOKENIZER is None:
-        try:
-            _TOKENIZER = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
-            if _TOKENIZER.pad_token_id is None:
-                if _TOKENIZER.eos_token_id is not None:
-                    _TOKENIZER.pad_token = _TOKENIZER.eos_token
-                    _TOKENIZER.pad_token_id = _TOKENIZER.eos_token_id
-                else:
-                    _TOKENIZER.add_special_tokens({"pad_token": "<|endoftext|>"})
-                    _TOKENIZER.pad_token_id = _TOKENIZER.convert_tokens_to_ids("<|endoftext|>")
-                logger.info(f"Set pad_token_id: {_TOKENIZER.pad_token_id}")
-            if _TOKENIZER.eos_token_id is None:
-                _TOKENIZER.eos_token = "<|endoftext|>"
-                _TOKENIZER.eos_token_id = _TOKENIZER.convert_tokens_to_ids("<|endoftext|>")
-                logger.info(f"Set eos_token_id: {_TOKENIZER.eos_token_id}")
-            logger.info("Tokenizer loaded successfully")
-        except Exception as e:
-            logger.error(f"Failed to load tokenizer: {type(e).__name__}: {e}")
-            raise
-    return _TOKENIZER
-
-def get_model():
-    global _MODEL
-    if _MODEL is None:
-        try:
-            logger.info("Loading model on CPU with FP16")
-            _MODEL = AutoModelForCausalLM.from_pretrained(
-                MODEL_NAME,
-                torch_dtype=torch.float16,
-                device_map="cpu",
-                low_cpu_mem_usage=True,
-            )
-            _MODEL.eval()
-            logger.info("Model loaded successfully")
-        except Exception as e:
-            logger.error(f"Failed to load model: {type(e).__name__}: {e}")
-            raise
-    return _MODEL
-
-# ------------ Prompt builder ------------
-def build_instructions(personality, level, topic):
-    return (
-        f"Plutus AI tutor for a {personality} learner at {level} level. "
-        f"Explain {topic} in a tone that fits learner's personality and tech level with examples. "
-        "Keep it 250–500 words. End with 'Takeaway:'."
-    )
-
-def build_model_input(tokenizer, personality, level, topic):
-    user_msg = build_instructions(personality, level, topic)
-    return f"System: You are a personalised Cardano Plutus tutor, your job is to make Plutus easy to learn based on different learner personalities, adapt your teaching style accordingly.\nUser: {user_msg}\nAssistant:"
-
-# ------------ Generation function ------------
-def generate(personality, level, topic, max_new_tokens=200):
-    try:
-        tokenizer = get_tokenizer()
-        model = get_model()
-        prompt = build_model_input(tokenizer, personality, level, topic)
-
-        device = torch.device("cpu")
-        logger.info("Generating on CPU")
-        inputs = tokenizer(prompt, return_tensors="pt")
-        input_len = inputs["input_ids"].shape[1]
-        inputs = {k: v.to(device) for k, v in inputs.items()}
-
-        start_time = time.time()
-        with torch.inference_mode():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=max_new_tokens,
-                temperature=0.3,
-                top_p=0.4,
-                do_sample=True,
-                repetition_penalty=1.3,
-                eos_token_id=tokenizer.eos_token_id,
-                pad_token_id=tokenizer.pad_token_id,
-                stop_strings=["Takeaway:"]
-            )
-        logger.info(f"Generation took {time.time() - start_time:.2f} seconds")
-
-        gen_ids = outputs[0][input_len:]
-        text = tokenizer.decode(gen_ids, skip_special_tokens=True).strip()
-        if not text:
-            text = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True).strip()
-
-        # Remove prompt if present
-        if prompt in text:
-            text = text.replace(prompt, "").strip()
-
-        # Truncate at Takeaway
-        if "Takeaway:" in text:
-            text = text[:text.index("Takeaway:") + len("Takeaway:") + 100].strip()
-        elif not text.endswith("Takeaway:"):
-            text += "\nTakeaway: (Summary not fully generated due to token limit)."
-
-        return text if text else "Generation failed. Try regenerating."
-    except Exception as e:
-        logger.error(f"Generation error: {type(e).__name__}: {e}")
-        return f"Error during generation: {str(e)}. Try regenerating or using a smaller model."
-
-# ------------ Orchestrator with retry logic ------------
-def orchestrator(personality, level, topic, max_retries=3):
-    if not personality or not level or not topic:
-        return "Select personality, expertise, and topic to get an explanation."
-    
-    logger.info("Yielding loading message")
-    yield "Generating response, please wait (~1–2 minutes on CPU)..."
-    
-    for attempt in range(max_retries):
-        try:
-            logger.info(f"Generation attempt {attempt + 1}/{max_retries}")
-            result = generate(personality, level, topic)
-            if result.startswith("Error during generation"):
-                raise Exception(result)
-            logger.info("Generation completed successfully")
-            return result
-        except Exception as e:
-            logger.error(f"Attempt {attempt + 1}/{max_retries} failed: {type(e).__name__}: {e}")
-            if attempt == max_retries - 1:
-                return (
-                    "Failed to generate after multiple attempts. "
-                    "Click **Regenerate** or try again later. "
-                    "If this persists, try a smaller model or check logs for errors."
-                )
-
-# ------------ Gradio UI ------------
-with gr.Blocks(theme="default") as iface:
-    gr.Markdown(
-        "## Cardano Plutus AI Assistant\n"
-        "Select **Learning Personality**, **Expertise Level**, and **Topic**, then click **Generate**. "
-        "Note: Generation may take ~1–2 minutes on CPU."
-    )
-
-    with gr.Row():
-        personality = gr.Dropdown(
-            choices=["Dyslexic", "Autistic", "Expressive"],
-            label="Learning Personality",
-            value=None,
-            allow_custom_value=False,
-            scale=1,
-        )
-        level = gr.Dropdown(
-            choices=["Beginner", "Intermediate", "Advanced"],
-            label="Expertise Level",
-            value=None,
-            allow_custom_value=False,
-            scale=1,
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+model.eval()
+
+if torch.cuda.is_available():
+    model.to("cuda")
+
+# Response function
+def generate_response(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=500,
+            temperature=0.3,
+            top_p=0.3,
+            do_sample=True,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
         )
-        topic = gr.Dropdown(
-            choices=[
-                "What is Plutus?",
-                "Why should I learn Plutus?",
-                "How does Plutus work?",
-                "What are the benefits of learning Plutus?",
-                "What are some common use cases for Plutus?",
-                "What are some resources for learning Plutus?",
-                "What are some best practices for using Plutus?",
-                "What are some common mistakes to avoid when using Plutus?",
-                "What are some advanced concepts in Plutus?",
-                "What are some future developments in Plutus?",
-            ],
-            label="Topic",
-            value=None,
-            allow_custom_value=False,
-            scale=2,
-        )
-
-    with gr.Row():
-        generate_btn = gr.Button("Generate")
-        regen = gr.Button("🔁 Regenerate")
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-    output = gr.Textbox(
-        label="Model Response",
-        lines=12,
-        interactive=False,
-        show_copy_button=True,
-        placeholder="Your tailored explanation will appear here…",
-    )
+    # Remove the prompt from the output to return only the answer
+    if response.startswith(prompt):
+        response = response[len(prompt):].strip()
 
-    generate_btn.click(orchestrator, [personality, level, topic], output, queue=True)
-    regen.click(orchestrator, [personality, level, topic], output, queue=True)
+    return response
 
-# Enable queue
-iface.queue()
+# Gradio UI
+demo = gr.Interface(
+    fn=generate_response,
+    inputs=gr.Textbox(label="Enter your prompt", lines=4, placeholder="Learn about Plutus..."),
+    outputs=gr.Textbox(label="Model Response"),
+    title="Cardano Plutus AI Assistant",
+    description="Your Personalised Plutus Tutor."
+)
 
-if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch()