Update app.py

app.py

@@ -1,54 +1,33 @@
 import gradio as gr
 import torch
-import logging, traceback
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    TextIteratorStreamer,
-    StoppingCriteria,
-    StoppingCriteriaList,
-)
+import logging
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
 from threading import Thread
 
-# ---------------- Logging ----------------
+# Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# ---------------- Model & Tokenizer ----------------
+# Load model & tokenizer
 MODEL_NAME = "ubiodee/Plutus_Tutor_new"
 
 try:
     logger.info("Loading tokenizer...")
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
-
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
     logger.info("Loading model...")
-    has_cuda = torch.cuda.is_available()
-    dtype = torch.float16 if has_cuda else torch.float32  # safer on CPU
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
         device_map="auto",
-        torch_dtype=dtype,
-        low_cpu_mem_usage=True,
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True
     )
     model.eval()
-
-    # Ensure pad/eos are sensible; if we add a token, resize embeddings
-    added = False
-    if tokenizer.pad_token_id is None:
-        if tokenizer.eos_token is not None:
-            tokenizer.pad_token = tokenizer.eos_token
-        else:
-            tokenizer.add_special_tokens({"pad_token": "</s>"})
-            added = True
-    if added:
-        model.resize_token_embeddings(len(tokenizer))
-
     logger.info("Model and tokenizer loaded successfully.")
 except Exception as e:
     logger.error(f"Error loading model or tokenizer: {str(e)}")
     raise
 
-# ---------------- UI Options ----------------
+# Define options for dropdowns
 PERSONALITY_TYPES = ["Autistic", "Dyslexic", "Expressive", "Nerd", "Visual", "Other"]
 PROGRAMMING_LEVELS = ["Beginner", "Intermediate", "Professional"]
 TOPICS = [
@@ -57,117 +36,89 @@ TOPICS = [
     "Smart Contracts",
     "Versioning in Plutus",
     "Monad",
-    "Other",
+    "Other"
 ]
 
-# ---------------- Prompting ----------------
-END_SENTINEL = "[END]"
-
+# Improved prompt template for autism-friendly, focused response
 def create_prompt(personality, level, topic):
-    return (
-        f"Explain {topic} in Plutus for a {level} programmer with {personality} traits. "
-        f"Use only basic words and clear examples. Use a physical object analogy (e.g., a lock or checklist) tied to {topic}. "
-        f"Avoid jargon like 'blockchain,' 'ledger,' 'Haskell,' 'decentralized,' 'cyber,' 'e-commerce,' 'formal verification,' or 'immutability.' "
-        f"Use short sentences (6-8 words). Use exactly 3 numbered points for key ideas. Each point must have 5-10 words. "
-        f"Bold the first word of each point. Structure the response: 2-sentence introduction, 3 numbered points, 1-sentence conclusion. "
-        f"For Autistic traits, use literal language, numbered lists, and **bold key terms**. Repeat key ideas for clarity. "
-        f"Avoid abstract terms unless concrete. Do not repeat the topic or prompt. Do not simulate a conversation, ask questions, or discuss unrelated topics. "
-        f"Use a direct, instructional tone without 'I' or 'we'. "
-        f"End with a summary sentence on {topic}'s importance, then write {END_SENTINEL} and nothing else."
-    )
-
-# ---------------- Stop on substring ----------------
-class StopOnSubstrings(StoppingCriteria):
-    def __init__(self, tokenizer, stop_strings):
-        self.stop_ids = [tokenizer.encode(s, add_special_tokens=False) for s in stop_strings]
-
-    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
-        for seq in self.stop_ids:
-            L = len(seq)
-            if L and input_ids.shape[1] >= L:
-                if torch.equal(input_ids[0, -L:], torch.tensor(seq, device=input_ids.device)):
-                    return True
-        return False
+    return f"Explain {topic} in Plutus for a {level} programmer with {personality} traits. Use only basic words and clear examples, like comparing validation to a lock. Avoid jargon like 'blockchain,' 'ledger,' 'Haskell,' 'decentralized,' 'formal verification,' or 'immutability.' Use short sentences (8 words or less). Use exactly 3 bullet points for key ideas. Each point must be under 15 words. Include one simple analogy. Structure the response: 2-sentence introduction, 3 bullet points, 1-sentence conclusion. For Autistic traits, use literal language, avoid abstract terms, and ensure a predictable format. Do not repeat the topic or prompt. Do not simulate a conversation, ask questions, or list unrelated terms. Use a direct, instructional tone without 'I,' 'we,' or conversational phrases. End with a summary sentence on the topic’s importance. Add extra line breaks between sections for readability."
 
-# ---------------- Generation (STREAMING) ----------------
+# Response function with improved parameters
 def generate_response(personality, level, topic):
     try:
         logger.info("Processing selections...")
         prompt = create_prompt(personality, level, topic)
         inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-        # Keep your original streaming pattern; avoid version-sensitive args
-        streamer = TextIteratorStreamer(
-            tokenizer,
-            skip_prompt=True,
-            skip_special_tokens=True,
-            # no timeout arg (some Gradio/HF versions don't support it)
-        )
-
-        stopping = StoppingCriteriaList([StopOnSubstrings(tokenizer, [END_SENTINEL])])
-
+        
+        # Use streamer for token-by-token generation
+        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        
         generation_kwargs = {
             **inputs,
             "streamer": streamer,
-            "max_new_tokens": 200,       # fits your format comfortably
-            "do_sample": False,          # deterministic to avoid tail babble
-            "no_repeat_ngram_size": 3,   # loop guard
-            "repetition_penalty": 1.1,   # mild anti-babble
-            "pad_token_id": tokenizer.pad_token_id,
-            "stopping_criteria": stopping,
-            "use_cache": True,
+            "max_new_tokens": 700,  # Increased to avoid cut-offs
+            "do_sample": True,
+            "temperature": 0.2,     # Lowered for more focused output
+            "top_p": 0.2,          # Lowered for more focused output
+            "repetition_penalty": 1.5,  # Increased to prevent repetition
+            "eos_token_id": tokenizer.eos_token_id,
+            "pad_token_id": tokenizer.pad_token_id
         }
-
-        # Only pass eos_token_id if it exists (avoid None issues)
-        if tokenizer.eos_token_id is not None:
-            generation_kwargs["eos_token_id"] = tokenizer.eos_token_id
-
-        thread = Thread(target=model.generate, kwargs=generation_kwargs, daemon=True)
+        
+        # Run generation in a separate thread
+        thread = Thread(target=model.generate, kwargs=generation_kwargs)
         thread.start()
-
+        
         generated_text = ""
         for new_text in streamer:
             generated_text += new_text
-
-            # Hard stop the moment we see the sentinel
-            if END_SENTINEL in generated_text:
-                yield generated_text.split(END_SENTINEL)[0].rstrip()
-                return
-
             yield generated_text.strip()
-
+        
         logger.info("Response generated successfully.")
-    except Exception:
-        err = traceback.format_exc()
-        logger.error(err)
-        # Show full traceback in UI for quick debugging
-        yield "Error:\n" + err
+    except Exception as e:
+        logger.error(f"Error during generation: {str(e)}")
+        yield f"Error: {str(e)}"
 
-# ---------------- Gradio UI ----------------
+# Gradio UI with dropdowns and button
 with gr.Blocks(title="Cardano Plutus AI Assistant") as demo:
     gr.Markdown("### Your Personalised Plutus Tutor")
     gr.Markdown("Select your personality type, programming level, and topic, then click Generate.")
-
-    personality = gr.Dropdown(choices=PERSONALITY_TYPES, label="Personality Type", value="Autistic")
-    level = gr.Dropdown(choices=PROGRAMMING_LEVELS, label="Programming Level", value="Beginner")
-    topic = gr.Dropdown(choices=TOPICS, label="Topic", value="Introduction to Validation")
-
+    
+    personality = gr.Dropdown(
+        choices=PERSONALITY_TYPES,
+        label="Personality Type",
+        value="Autistic"
+    )
+    level = gr.Dropdown(
+        choices=PROGRAMMING_LEVELS,
+        label="Programming Level",
+        value="Beginner"
+    )
+    topic = gr.Dropdown(
+        choices=TOPICS,
+        label="Topic",
+        value="Introduction to Validation"
+    )
+    
     generate_btn = gr.Button("Generate")
-
+    
     output = gr.Textbox(
         label="Model Response",
         show_label=True,
         lines=10,
-        placeholder="Generated content will appear here...",
+        placeholder="Generated content will appear here..."
     )
-
+    
     generate_btn.click(
         fn=generate_response,
         inputs=[personality, level, topic],
-        outputs=output,
+        outputs=output
     )
 
-logger.info("Launching Gradio interface...")
-# Keep it version-agnostic: enable queueing without extra args
-demo.queue()
-demo.launch()
+# Launch the app
+try:
+    logger.info("Launching Gradio interface...")
+    demo.launch()
+except Exception as e:
+    logger.error(f"Error launching Gradio: {str(e)}")
+    raise