Update inference.py

inference.py

@@ -102,7 +102,7 @@ def _apply_repetition_penalty(
     logits: torch.Tensor,
     generated_ids: list[int],
     penalty: float,
-    prompt_id_set: "list[int] | set[int] | None" = None,
+    prompt_counter: "Counter[int] | None" = None,
     max_count: int = 6,
 ) -> torch.Tensor:
     """CTRL-style penalty, COUNT-SCALED: a token seen N times in the
@@ -136,12 +136,12 @@ def _apply_repetition_penalty(
     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
+    interference. prompt_counter (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
-    if prompt_id_set:
+    if prompt_counter:
         # 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 the prompt -- tokens that the
@@ -149,9 +149,8 @@ def _apply_repetition_penalty(
         # 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 = _build_prompt_remaining(
-            list(prompt_id_set) if isinstance(prompt_id_set, set) else list(prompt_id_set)
-        )
+        # Copy the pre-built Counter (O(unique tokens), not O(prompt_len)).
+        prompt_remaining = prompt_counter.copy()
         for tid in generated_ids:
             if prompt_remaining.get(tid, 0) > 0:
                 prompt_remaining[tid] -= 1
@@ -224,13 +223,6 @@ def _top_p_filter(logits: torch.Tensor, top_p: float) -> torch.Tensor:
     return torch.full_like(logits, float("-inf")).scatter(1, sorted_idx, sorted_logits)
 
 
-def _build_prompt_remaining(prompt_ids: list[int]) -> Counter[int]:
-    """Return the prompt token multiplicities needed by the repetition
-    penalty walker. A plain set hides duplicate counts, which underestimates
-    how many times a repeated prompt token should be available to the counter.
-    """
-    return Counter(prompt_ids)
-
 
 def _clip_max_new_tokens(max_new_tokens: int, prompt_ids: list[int], block_size: int) -> int:
     """Prevent the generation loop from silently overwriting the prompt by
@@ -273,7 +265,10 @@ def generate(
 
     # Built once from the original prompt and kept as a multiset so
     # repetition penalty excludes prompt tokens by duplicate-aware count.
-    prompt_id_set: list[int] = list(prompt_ids)
+    prompt_ids_reference: list[int] = list(prompt_ids)
+    # Counter built once here; _apply_repetition_penalty copies it per call
+    # (O(unique tokens)) instead of rebuilding from the list (O(prompt_len)).
+    prompt_counter: Counter[int] = Counter(prompt_ids_reference)
 
     # Clip the requested budget to the remaining context budget so the
     # context window never drops the prompt silently when max_new_tokens
@@ -305,7 +300,7 @@ def generate(
             if finished[b]:
                 continue
             row = logits[b:b + 1]
-            row = _apply_repetition_penalty(row, generated_ids[b], repetition_penalty, prompt_id_set=prompt_id_set)
+            row = _apply_repetition_penalty(row, generated_ids[b], repetition_penalty, prompt_counter=prompt_counter)
             row = _block_repeated_ngrams(row, generated_ids[b], no_repeat_ngram_size)
             row = row / max(temperature, 1e-5)
             if top_k is not None:
@@ -314,6 +309,12 @@ def generate(
                 row = _top_p_filter(row, top_p)
             logits[b:b + 1] = row
 
+        # Guard: if any row is all -inf (over-aggressive filtering), reset it
+        # to uniform so multinomial doesn't receive NaN probabilities.
+        all_neg_inf = ~torch.isfinite(logits).any(dim=-1)  # [B]
+        if all_neg_inf.any():
+            logits[all_neg_inf] = 0.0  # uniform before softmax
+
         probs = torch.softmax(logits, dim=-1)
         next_ids = torch.multinomial(probs, num_samples=1)  # (num_samples, 1)
 
@@ -324,15 +325,22 @@ def generate(
                 next_ids[b, 0] = idx[b, -1]
                 continue
             tid = next_ids[b, 0].item()
-            generated_ids[b].append(tid)
             if eos_id is not None and tid == eos_id:
                 finished[b] = True
+            else:
+                generated_ids[b].append(tid)
 
         idx = torch.cat([idx, next_ids], dim=1)
 
         if stream and num_samples == 1:
             new_text = tokenizer.decode(generated_ids[0])
-            print(new_text[len(prev_text[0]):], end="", flush=True)
+            if new_text.startswith(prev_text[0]):
+                print(new_text[len(prev_text[0]):], end="", flush=True)
+            else:
+                # BPE decode is not prefix-stable (a new token can change how
+                # earlier bytes are decoded, e.g. completing a multi-byte char).
+                # Reprint the whole line to stay correct.
+                print("\r" + new_text, end="", flush=True)
             prev_text[0] = new_text
 
     if stream and num_samples == 1: