Update app.py

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 import torch
-import logging
+import logging, traceback
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
@@ -22,18 +22,27 @@ try:
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
 
     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",
-        dtype=torch.float16,          # transformers now prefers `dtype` over `torch_dtype`
+        torch_dtype=dtype,
         low_cpu_mem_usage=True,
-    ).eval()
+    )
+    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_id is not 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)}")
@@ -67,7 +76,7 @@ def create_prompt(personality, level, topic):
         f"End with a summary sentence on {topic}'s importance, then write {END_SENTINEL} and nothing else."
     )
 
-# ---------------- Stopping on substring ----------------
+# ---------------- 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]
@@ -87,11 +96,12 @@ def generate_response(personality, level, topic):
         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,
-            timeout=0.02,  # flush small chunks quickly
+            # no timeout arg (some Gradio/HF versions don't support it)
         )
 
         stopping = StoppingCriteriaList([StopOnSubstrings(tokenizer, [END_SENTINEL])])
@@ -99,31 +109,39 @@ def generate_response(personality, level, topic):
         generation_kwargs = {
             **inputs,
             "streamer": streamer,
-            "max_new_tokens": 200,
-            "do_sample": False,
-            "no_repeat_ngram_size": 3,
-            "repetition_penalty": 1.1,
-            "eos_token_id": tokenizer.eos_token_id,
+            "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,
         }
 
+        # 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)
         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 as e:
-        logger.error(f"Error during generation: {str(e)}")
-        yield f"Error: {str(e)}"
+    except Exception:
+        err = traceback.format_exc()
+        logger.error(err)
+        # Show full traceback in UI for quick debugging
+        yield "Error:\n" + err
 
 # ---------------- Gradio UI ----------------
 with gr.Blocks(title="Cardano Plutus AI Assistant") as demo:
@@ -143,16 +161,13 @@ with gr.Blocks(title="Cardano Plutus AI Assistant") as demo:
         placeholder="Generated content will appear here...",
     )
 
-    # (Optional) Per-event concurrency control in Gradio 4+
     generate_btn.click(
         fn=generate_response,
         inputs=[personality, level, topic],
         outputs=output,
-        concurrency_limit=1,      # <- replaces old global concurrency_count
     )
 
 logger.info("Launching Gradio interface...")
-
-# ✅ Gradio 4+ queue config (no more `concurrency_count`)
-demo.queue(default_concurrency_limit=1, max_size=20)
+# Keep it version-agnostic: enable queueing without extra args
+demo.queue()
 demo.launch()