Update app.py

app.py

@@ -33,15 +33,21 @@ for name in MODELS:
 
 
 # ---- manual sampling fallback (works for ANY model, ignores custom generate) ----
-def manual_sample(model, input_ids, max_new_tokens, temperature, top_k):
+def manual_sample(model, input_ids, max_new_tokens, temperature, top_k, min_new_tokens=0):
     """Generic top-k sampler. Bypasses any custom .generate() method."""
     generated = input_ids.clone()
-    for _ in range(int(max_new_tokens)):
+    eos_id = getattr(model.config, "eos_token_id", None)
+
+    for i in range(int(max_new_tokens)):
         with torch.no_grad():
             outputs = model(generated)
             logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
         next_logits = logits[0, -1, :] / max(temperature, 1e-5)
 
+        # Suppress EOS for the first `min_new_tokens` steps
+        if i < min_new_tokens and eos_id is not None:
+            next_logits[eos_id] = float("-inf")
+
         if top_k > 0:
             k = min(int(top_k), next_logits.size(-1))
             top_vals, top_idx = torch.topk(next_logits, k)
@@ -54,8 +60,9 @@ def manual_sample(model, input_ids, max_new_tokens, temperature, top_k):
 
         generated = torch.cat([generated, next_token.view(1, 1)], dim=-1)
 
-        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:
+        # Only allow EOS to stop generation after min_new_tokens
+        if i >= min_new_tokens and eos_id is not None:
+            if next_token.item() == eos_id:
                 break
 
     return generated
@@ -98,6 +105,7 @@ def run_generate(model_name, prompt, max_new_tokens, temperature, top_k):
             out = model.generate(
                 input_ids,
                 max_new_tokens=int(max_new_tokens),
+                min_new_tokens=10,          # <-- force at least 10 new tokens
                 temperature=float(temperature),
                 top_k=int(top_k),
                 do_sample=True,
@@ -105,7 +113,10 @@ def run_generate(model_name, prompt, max_new_tokens, temperature, top_k):
         return decode_output(tokenizer, out)
     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)
+        out = manual_sample(
+            model, input_ids, max_new_tokens, temperature, top_k,
+            min_new_tokens=10,              # <-- same here
+        )
         return decode_output(tokenizer, out)