Update app.py

app.py

@@ -6,41 +6,42 @@ from transformers import (
     AutoTokenizer,
     TextIteratorStreamer,
     StoppingCriteria,
-    StoppingCriteriaList
+    StoppingCriteriaList,
 )
 from threading import Thread
 
+# ---------------- Logging ----------------
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# ---------------- Model & Tokenizer ----------------
 MODEL_NAME = "ubiodee/Plutus_Tutor_new"
 
-logger.info("Loading tokenizer...")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-logger.info("Loading model...")
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    device_map="auto",
-    torch_dtype=torch.float16,
-    low_cpu_mem_usage=True
-).eval()
-
-# Ensure pad/eos set sensibly
-if tokenizer.pad_token_id is None:
-    tokenizer.pad_token = tokenizer.eos_token or tokenizer.pad_token or "</s>"
-
-def eos_id_candidates(tok):
-    ids = set()
-    for tok_str in ["</s>", "<|eot_id|>", "<|end|>", "<|im_end|>"]:
-        tid = tok.convert_tokens_to_ids(tok_str)
-        if tid is not None and tid != -1:
-            ids.add(tid)
-    if tok.eos_token_id is not None:
-        ids.add(tok.eos_token_id)
-    return list(ids) if ids else None
-
-EOS_IDS = eos_id_candidates(tokenizer)
+try:
+    logger.info("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
 
+    logger.info("Loading model...")
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map="auto",
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True,
+    )
+    model.eval()
+
+    # Make sure pad/eos are sensible to avoid warnings/crashes
+    if tokenizer.pad_token_id is None:
+        if tokenizer.eos_token_id is not None:
+            tokenizer.pad_token = tokenizer.eos_token
+        else:
+            tokenizer.add_special_tokens({"pad_token": "</s>"})
+    logger.info("Model and tokenizer loaded successfully.")
+except Exception as e:
+    logger.error(f"Error loading model or tokenizer: {str(e)}")
+    raise
+
+# ---------------- UI Options ----------------
 PERSONALITY_TYPES = ["Autistic", "Dyslexic", "Expressive", "Nerd", "Visual", "Other"]
 PROGRAMMING_LEVELS = ["Beginner", "Intermediate", "Professional"]
 TOPICS = [
@@ -49,66 +50,79 @@ TOPICS = [
     "Smart Contracts",
     "Versioning in Plutus",
     "Monad",
-    "Other"
+    "Other",
 ]
 
+# ---------------- Prompting ----------------
 END_SENTINEL = "[END]"
 
 def create_prompt(personality, level, topic):
+    # Keep your structure & tone, add explicit end signal
     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 tied to {topic}. "
-        f"Avoid jargon like 'blockchain,' 'ledger,' 'Haskell,' 'decentralized,' 'cyber,' 'e-commerce,' "
-        f"'formal verification,' or 'immutability.' Use short sentences (6-8 words). "
-        f"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, "
-        f"3 numbered points, 1-sentence conclusion. For Autistic traits, use literal language, "
-        f"numbered lists, and **bold key terms**. Repeat key ideas for clarity. "
-        f"Avoid abstract terms unless concrete. Do not repeat the topic or prompt. "
-        f"Do not ask questions. Use a direct, instructional tone. "
-        f"End with a summary sentence on {topic}’s importance, then write {END_SENTINEL} and nothing else."
+        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."
     )
 
-# StoppingCriteria that halts when a stop substring appears
+# ---------------- Stopping on substring ----------------
 class StopOnSubstrings(StoppingCriteria):
     def __init__(self, tokenizer, stop_strings):
+        self.tokenizer = tokenizer
+        # Pre-tokenize stop strings for fast suffix checks
         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:
+        # Stop if any stop_ids match the suffix of the generated sequence
         for seq in self.stop_ids:
             L = len(seq)
-            if L == 0: 
+            if L == 0:
                 continue
             if input_ids.shape[1] >= L:
-                if torch.equal(input_ids[0, -L:], torch.tensor(seq, device=input_ids.device)):
+                if torch.equal(
+                    input_ids[0, -L:],
+                    torch.tensor(seq, device=input_ids.device),
+                ):
                     return True
         return False
 
+# ---------------- Generation (STREAMING) ----------------
 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)
 
-        streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+        # Keep streamer + background thread approach (as in your working version)
+        streamer = TextIteratorStreamer(
+            tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True,
+            timeout=0.02,  # flush small chunks quickly
+        )
 
         stopping = StoppingCriteriaList([StopOnSubstrings(tokenizer, [END_SENTINEL])])
 
+        # Tighter, deterministic decoding to avoid trailing garbage
         generation_kwargs = {
             **inputs,
             "streamer": streamer,
-
-            # <<< Tighter, safer defaults for short, structured outputs >>>
-            "max_new_tokens": 180,
-            "do_sample": False,                # deterministic; avoids tail babble
-            "no_repeat_ngram_size": 3,
-            "repetition_penalty": 1.1,
-            "eos_token_id": EOS_IDS,           # list of possible EOS tokens
+            "max_new_tokens": 200,         # your format fits well under this
+            "do_sample": False,            # deterministic; helps finish cleanly
+            "no_repeat_ngram_size": 3,     # avoid loops
+            "repetition_penalty": 1.1,     # gentle anti-babble
+            "eos_token_id": tokenizer.eos_token_id,
             "pad_token_id": tokenizer.pad_token_id,
             "stopping_criteria": stopping,
             "use_cache": True,
         }
 
+        # Run generation in a separate thread so we can iterate the streamer
         thread = Thread(target=model.generate, kwargs=generation_kwargs, daemon=True)
         thread.start()
 
@@ -116,12 +130,12 @@ def generate_response(personality, level, topic):
         for new_text in streamer:
             generated_text += new_text
 
-            # Defensive: cut off sentinel if it appears in the stream
+            # Hard stop if sentinel appears; strip it from output
             if END_SENTINEL in generated_text:
-                clean = generated_text.split(END_SENTINEL)[0].rstrip()
-                yield clean
+                yield generated_text.split(END_SENTINEL)[0].rstrip()
                 return
 
+            # Stream progressively (exactly like your earlier working version)
             yield generated_text.strip()
 
         logger.info("Response generated successfully.")
@@ -129,6 +143,7 @@ def generate_response(personality, level, topic):
         logger.error(f"Error during generation: {str(e)}")
         yield f"Error: {str(e)}"
 
+# ---------------- Gradio UI ----------------
 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.")
@@ -138,9 +153,21 @@ with gr.Blocks(title="Cardano Plutus AI Assistant") as demo:
     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...")
 
-    generate_btn.click(fn=generate_response, inputs=[personality, level, topic], outputs=output)
+    output = gr.Textbox(
+        label="Model Response",
+        show_label=True,
+        lines=10,
+        placeholder="Generated content will appear here...",
+    )
+
+    generate_btn.click(
+        fn=generate_response,
+        inputs=[personality, level, topic],
+        outputs=output,
+    )
 
+# Ensure true streaming in Gradio
 logger.info("Launching Gradio interface...")
+demo.queue(concurrency_count=1, max_size=20)
 demo.launch()