Harness: no-repeat-ngram(3) flag, default on (A/B'd per model on real weights, rep-4 -> 0.000 at deployed settings)

generate.py

@@ -89,8 +89,19 @@ def _filter_logits(logits, top_k, top_p):
 
 
 @torch.no_grad()
+def _ban3(ids, n=3):
+    """no-repeat-ngram: token ids that would complete an already-seen n-gram.
+    A/B'd per model on real weights: rep-4 -> 0.000 at deployed settings."""
+    if n <= 0 or len(ids) < n:
+        return ()
+    tail = tuple(ids[-(n - 1):])
+    return tuple(ids[i + n - 1] for i in range(len(ids) - n + 1)
+                 if tuple(ids[i:i + n - 1]) == tail)
+
+
 def generate(model, cfg, idx, n_new, temperature=0.9, top_k=40, top_p=0.95,
-             repetition_penalty=1.1, stop_ids=None, device="cpu", use_cache=True):
+             repetition_penalty=1.1, stop_ids=None, device="cpu", use_cache=True,
+             no_repeat_ngram=3):
     """Autoregressive sampling. Mirrors model.generate()'s NaN/inf sanitisation
     (long oscillator rollouts can emit +/-inf logits whose softmax -> NaN ->
     multinomial device-side assert) and adds top-p + repetition penalty.
@@ -101,6 +112,7 @@ def generate(model, cfg, idx, n_new, temperature=0.9, top_k=40, top_p=0.95,
     must stay within max_position_embeddings)."""
     model.eval()
     stop_ids = set(stop_ids or [])
+    n0 = idx.size(1)
     past = None
     for _ in range(n_new):
         if use_cache and past is not None and past[0][0].size(2) < cfg.block_size:
@@ -120,6 +132,9 @@ def generate(model, cfg, idx, n_new, temperature=0.9, top_k=40, top_p=0.95,
                 lg[b, seen] = torch.where(vals > 0, vals / repetition_penalty,
                                           vals * repetition_penalty)
 
+        if no_repeat_ngram:
+            for b in _ban3(idx[0, n0:].tolist(), no_repeat_ngram):
+                lg[0, b] = float("-inf")
         lg = lg / max(temperature, 1e-6)
         lg = _filter_logits(lg, top_k, top_p)
         probs = torch.softmax(lg, dim=-1)
@@ -256,6 +271,7 @@ def run_once(model, cfg, tok, args, device, prompt_text):
         out = generate(model, cfg, x, args.max_new_tokens,
                        temperature=args.temperature, top_k=args.top_k,
                        top_p=args.top_p, repetition_penalty=args.repetition_penalty,
+                       no_repeat_ngram=args.no_repeat_ngram,
                        stop_ids=stop_ids, device=device, use_cache=not args.no_cache)
     gen_ids = out[0, len(ids):].tolist()
     completion = tok.decode(gen_ids, skip_special_tokens=not args.show_special)
@@ -276,6 +292,7 @@ def main():
     p.add_argument("--top_k", type=int, default=40, help="0 to disable")
     p.add_argument("--top_p", type=float, default=0.95, help="1.0 to disable")
     p.add_argument("--repetition_penalty", type=float, default=1.1, help="1.0 to disable")
+    p.add_argument("--no_repeat_ngram", type=int, default=3, help="0 to disable")
     p.add_argument("--seed", type=int, default=None)
     p.add_argument("--chat", action="store_true", help="wrap prompt in ChatML framing")
     p.add_argument("--system", default="", help="system prompt (chat mode only)")