Update app.py

app.py

@@ -12,20 +12,14 @@ MODELS = {
 loaded = {}
 load_errors = {}
 
+
 def load_model(name):
     if name in loaded or name in load_errors:
         return
     repo = MODELS[name]
     try:
-        # 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
-        )
+        tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
+        model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True)
         model.eval()
         loaded[name] = (model, tokenizer)
         print(f"[loaded] {name}")
@@ -33,6 +27,7 @@ def load_model(name):
         load_errors[name] = f"{type(e).__name__}: {e}"
         print(f"[load failed] {name}: {e}")
 
+
 for name in MODELS:
     load_model(name)
 
@@ -67,7 +62,7 @@ def manual_sample(model, input_ids, max_new_tokens, temperature, top_k):
 
 
 def encode_prompt(tokenizer, prompt):
-    # FIX 2: more robust handling for both fast and slow tokenizers.
+    """Handle both fast and slow tokenizers, and flatten nested outputs."""
     if hasattr(tokenizer, "encode"):
         ids = tokenizer.encode(prompt)
         if hasattr(ids, "ids"):
@@ -75,7 +70,6 @@ def encode_prompt(tokenizer, prompt):
     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):
@@ -85,7 +79,7 @@ def encode_prompt(tokenizer, prompt):
 
 
 def decode_output(tokenizer, ids):
-    # FIX 3: normalize to a flat python list of ints before decoding.
+    """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):
@@ -98,8 +92,7 @@ def run_generate(model_name, prompt, max_new_tokens, temperature, top_k):
     model, tokenizer = loaded[model_name]
     input_ids = encode_prompt(tokenizer, prompt)
 
-    # 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 the model's own generate first, fall back to manual sampling on failure
     try:
         with torch.no_grad():
             out = model.generate(