Update inference.py

inference.py

@@ -97,28 +97,77 @@ def load_model(checkpoint_path: str, device: torch.device) -> tuple[GPT, GPTConf
 # Sampling
 # ─────────────────────────────────────────────────────────────────────────────
 
-def _apply_repetition_penalty(logits: torch.Tensor, generated_ids: list[int], penalty: float) -> torch.Tensor:
-    """CTRL-style penalty, COUNT-SCALED: a token seen N times gets divided
-    by penalty**N (if positive) / multiplied by penalty**N (if negative),
-    not a single flat divide regardless of occurrence count.
-
-    The original flat version (divide once no matter how many times a
-    token has recurred) is too weak on small/repetitive-domain models
-    (e.g. TinyStories-style data where "so so so happy!" is a genuinely
-    common training pattern): once the model locks onto a token, the
-    penalty never gets any stronger no matter how many times it fires,
-    so a token that's only mildly favored over the runner-up never gets
-    pushed low enough to escape the loop. Scaling by occurrence count
-    means each additional repeat compounds the penalty, so a loop that's
-    3-4 tokens deep gets meaningfully suppressed even if a single
-    application wouldn't have been enough.
+def _apply_repetition_penalty(
+    logits: torch.Tensor,
+    generated_ids: list[int],
+    penalty: float,
+    prompt_id_set: "set[int] | None" = None,
+    max_count: int = 6,
+) -> torch.Tensor:
+    """CTRL-style penalty, COUNT-SCALED: a token seen N times in the
+    *generated* portion gets divided by penalty**N (if positive) /
+    multiplied by penalty**N (if negative).
+
+    Count-scaling vs flat penalty
+    ------------------------------
+    The flat version (divide once regardless of occurrence count) is too
+    weak on small/repetitive-domain models (e.g. TinyStories-style data
+    where "so so so happy!" is a genuinely common training pattern): once
+    the model locks onto a token, the penalty never gets any stronger, so
+    a mildly favoured looping token never gets pushed low enough to escape.
+    Scaling by occurrence count means each additional repeat compounds the
+    penalty, so a 3-4 token-deep loop gets meaningfully suppressed even if
+    a single application wouldn't have been enough.
+
+    Why penalty counts are clamped to max_count
+    --------------------------------------------
+    Without a cap, a legitimate word that happens to appear many times in
+    a long generation (e.g. "the", "a", punctuation) accumulates a factor
+    of penalty**N that grows without bound and can suppress even
+    high-confidence logits into noise. max_count=6 lets count-scaling do
+    its job on real loops (which repeat in short runs) while preventing
+    runaway suppression of common tokens over a long context.
+
+    Why prompt tokens are excluded from the count
+    -----------------------------------------------
+    generated_ids is seeded with the full prompt before the first token is
+    produced. Without the exclusion, common tokens in a 20-token prompt
+    already arrive at the first generation step with count >= 2-3, giving
+    them a standing penalty of penalty**3 before the model has emitted a
+    single token. That's not repetition suppression -- it's prompt
+    interference. prompt_id_set (built once per generate() call from the
+    original prompt_ids) lets us count only within the generated portion.
     """
     if penalty == 1.0 or not generated_ids:
         return logits
     from collections import Counter
-    counts = Counter(generated_ids)
-    seen = torch.tensor(sorted(counts.keys()), device=logits.device, dtype=torch.long)
-    reps = torch.tensor([counts[t] for t in sorted(counts.keys())], device=logits.device, dtype=logits.dtype)
+    if prompt_id_set:
+        # Count only tokens that appear in the generated suffix (after the
+        # prompt). generated_ids includes the prompt prefix, so we walk the
+        # list and tally only tokens not in prompt_id_set -- tokens that the
+        # model itself produced. Prompt tokens still receive the flat
+        # penalty**1 if they recur in the generated portion (they're not
+        # immune; we just don't pre-charge them for being in the prompt).
+        gen_counts: Counter[int] = Counter()
+        prompt_remaining = Counter(prompt_id_set)
+        for tid in generated_ids:
+            if prompt_remaining.get(tid, 0) > 0:
+                prompt_remaining[tid] -= 1
+                continue
+            gen_counts[tid] += 1
+        counts = gen_counts
+    else:
+        counts = Counter(generated_ids)
+
+    if not counts:
+        return logits
+
+    sorted_tokens = sorted(counts.keys())
+    seen = torch.tensor(sorted_tokens, device=logits.device, dtype=torch.long)
+    # Clamp counts to max_count so common tokens don't accumulate a
+    # runaway factor over a long generation.
+    raw_reps = [min(counts[t], max_count) for t in sorted_tokens]
+    reps = torch.tensor(raw_reps, device=logits.device, dtype=logits.dtype)
     vals = logits[0, seen]
     factor = penalty ** reps
     logits[0, seen] = torch.where(vals > 0, vals / factor, vals * factor)
@@ -197,6 +246,10 @@ def generate(
     block_size = model.config.block_size
     eos_id = getattr(tokenizer, "eos_id", None)
 
+    # Built once from the original prompt so _apply_repetition_penalty can
+    # exclude prompt tokens from the count -- see that function's docstring.
+    prompt_id_set: set[int] = set(prompt_ids)
+
     generated_ids = [list(prompt_ids) for _ in range(num_samples)]
     finished = [False] * num_samples
     # Decode the FULL sequence each step and print only the new suffix,
@@ -222,7 +275,7 @@ def generate(
             if finished[b]:
                 continue
             row = logits[b:b + 1]
-            row = _apply_repetition_penalty(row, generated_ids[b], repetition_penalty)
+            row = _apply_repetition_penalty(row, generated_ids[b], repetition_penalty, prompt_id_set=prompt_id_set)
             row = _block_repeated_ngrams(row, generated_ids[b], no_repeat_ngram_size)
             row = row / max(temperature, 1e-5)
             if top_k is not None: