Update app.py

app.py

@@ -17,8 +17,15 @@ def load_model(name):
         return
     repo = MODELS[name]
     try:
-        tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
-        model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
+        # FIX 1: use_fast=False forces the slow GPT-2 tokenizer path,
+        # which reads vocab.json + merges.txt directly and never touches
+        # the broken tokenizer.json -> avoids the "missing field trim_offsets" error.
+        tokenizer = AutoTokenizer.from_pretrained(
+            repo, trust_remote_code=True, use_fast=False
+        )
+        model = AutoModelForCausalLM.from_pretrained(
+            repo, trust_remote_code=True
+        )
         model.eval()
         loaded[name] = (model, tokenizer)
         print(f"[loaded] {name}")
@@ -29,6 +36,7 @@ def load_model(name):
 for name in MODELS:
     load_model(name)
 
+
 # ---- manual sampling fallback (works for ANY model, ignores custom generate) ----
 def manual_sample(model, input_ids, max_new_tokens, temperature, top_k):
     """Generic top-k sampler. Bypasses any custom .generate() method."""
@@ -39,7 +47,6 @@ def manual_sample(model, input_ids, max_new_tokens, temperature, top_k):
             logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
         next_logits = logits[0, -1, :] / max(temperature, 1e-5)
 
-        # top-k filtering
         if top_k > 0:
             k = min(int(top_k), next_logits.size(-1))
             top_vals, top_idx = torch.topk(next_logits, k)
@@ -52,34 +59,47 @@ def manual_sample(model, input_ids, max_new_tokens, temperature, top_k):
 
         generated = torch.cat([generated, next_token.view(1, 1)], dim=-1)
 
-        # stop if we hit an EOS
         if hasattr(model.config, "eos_token_id") and model.config.eos_token_id is not None:
             if next_token.item() == model.config.eos_token_id:
                 break
 
     return generated
 
+
 def encode_prompt(tokenizer, prompt):
+    # FIX 2: more robust handling for both fast and slow tokenizers.
     if hasattr(tokenizer, "encode"):
         ids = tokenizer.encode(prompt)
         if hasattr(ids, "ids"):
             ids = ids.ids
     else:
         ids = tokenizer(prompt)["input_ids"]
+
+    # Ensure we always have a flat list of ints
+    if torch.is_tensor(ids):
+        ids = ids.tolist()
+    if isinstance(ids, list) and len(ids) > 0 and isinstance(ids[0], list):
+        ids = ids[0]
+
     return torch.tensor([ids], dtype=torch.long)
 
+
 def decode_output(tokenizer, ids):
-    try:
-        return tokenizer.decode(ids[0].tolist() if torch.is_tensor(ids[0]) else ids[0], skip_special_tokens=True)
-    except TypeError:
-        return tokenizer.decode(ids[0], skip_special_tokens=True)
+    # FIX 3: normalize to a flat python list of ints before decoding.
+    if torch.is_tensor(ids):
+        ids = ids.tolist()
+    if isinstance(ids, list) and len(ids) > 0 and isinstance(ids[0], list):
+        ids = ids[0]
+    return tokenizer.decode(ids, skip_special_tokens=True)
+
 
 @spaces.GPU(duration=30)
 def run_generate(model_name, prompt, max_new_tokens, temperature, top_k):
     model, tokenizer = loaded[model_name]
     input_ids = encode_prompt(tokenizer, prompt)
 
-    # try the model's own generate first, fall back to manual sampling
+    # FIX 4: some tiny custom models don't implement .generate() correctly,
+    # so we try it first and fall back to manual sampling on ANY exception.
     try:
         with torch.no_grad():
             out = model.generate(
@@ -90,11 +110,12 @@ def run_generate(model_name, prompt, max_new_tokens, temperature, top_k):
                 do_sample=True,
             )
         return decode_output(tokenizer, out)
-    except (TypeError, AttributeError, ValueError) as e:
+    except Exception as e:
         print(f"[fallback] {model_name} custom generate failed ({e}), using manual sampling")
         out = manual_sample(model, input_ids, max_new_tokens, temperature, top_k)
         return decode_output(tokenizer, out)
 
+
 def compare_all(prompt, max_new_tokens, temperature, top_k):
     results = []
     for name in MODELS:
@@ -111,6 +132,7 @@ def compare_all(prompt, max_new_tokens, temperature, top_k):
         results.append(text)
     return results
 
+
 with gr.Blocks(title="TinyBuddy Trilogy") as demo:
     gr.Markdown(
         """