Break 8B JSON runaway loops with a retry-only repetition penalty

Some posters drive the stock Llama-3.1-8B into a repetition loop that
re-emits the creators array and never closes the JSON, grinding to the
18k token cap and producing an unparseable [ERR] result (e.g. 16083265).

Fix, without touching the healthy path:
- _generate() now returns (text, hit_eos); hit_eos is False when a pass
  reaches the token cap without emitting EOS -- the runaway signature.
- extract_json_with_retry() retries on that signal in addition to parse
  errors. This matters because the robust parser can salvage a truncated
  runaway into a parseable-but-gutted object (no error key) that would
  otherwise skip the retry.
- The retry/fallback passes apply repetition_penalty=1.15 to break the
  loop. It is NOT applied on the primary pass, so healthy posters (which
  emit EOS on pass 1 and never retry) stay byte-identical. 1.15 was tuned
  empirically -- 1.3 corrupts the verbatim transcript.

16083265 end-to-end: [ERR] (0 sections) -> PASS (16/16 sections, rougeL
0.86); healthy control 10890106 unchanged. Doc updated.

Author: jimnoneill

llama_generation_settings.md

@@ -12,10 +12,11 @@ All settings live in `poster2json/extract.py`. Code anchors are given per settin
 
 | Setting | Value | Where |
 |---------|-------|-------|
-| Decoding | greedy (`do_sample=False`) | `_generate()`, line 1357 |
+| Decoding | greedy (`do_sample=False`) | `_generate()` |
 | Sampling params | none (no `temperature` / `top_p` / `top_k` / `num_beams`) | `_generate()` |
 | Primary output budget | `MAX_JSON_TOKENS = 18000` | line 46 |
 | Retry/fallback budget | `MAX_RETRY_TOKENS = 24000` | line 47 |
+| Retry repetition penalty | `RETRY_REPETITION_PENALTY = 1.15` (retry/fallback only) | line 50 |
 | Input gate | `MAX_INPUT_TOKENS = 15000` | line 48 |
 | EOS handling | `_JsonBraceProcessor` (custom `LogitsProcessor`) | line 1282 |
 | Default quantization | 4-bit NF4 (`bnb_4bit_compute_dtype=bfloat16`, double-quant) | `load_json_model()`, line 1227 |
@@ -62,14 +63,18 @@ The model's native context is 128K, so 15k-in + 24k-out sits comfortably inside
 
 ### Retry ladder
 
-`extract_json_with_retry()` (line 2251) escalates only when the parse fails or the output looks truncated:
+`extract_json_with_retry()` escalates when the parse fails, the output looks truncated, **or the pass ran to the token cap without emitting EOS** (`hit_eos == False`):
 
-1. **Primary** — full `EXTRACTION_PROMPT` @ `MAX_JSON_TOKENS` (18k).
-2. **Retry** — same full prompt @ `MAX_RETRY_TOKENS` (24k), if step 1 errored or truncated.
-3. **Fallback** — shorter `FALLBACK_PROMPT` @ `MAX_RETRY_TOKENS` (24k), if step 2 still errored or truncated.
+1. **Primary** — full `EXTRACTION_PROMPT` @ `MAX_JSON_TOKENS` (18k), plain greedy (no penalty).
+2. **Retry** — same full prompt @ `MAX_RETRY_TOKENS` (24k) with `repetition_penalty = RETRY_REPETITION_PENALTY`, if step 1 errored, was truncated, or hit the cap without EOS.
+3. **Fallback** — shorter `FALLBACK_PROMPT` @ `MAX_RETRY_TOKENS` (24k), same penalty, if step 2 still failed those checks.
 
 Each step is followed by `_robust_json_parse()` (hand-rolled repair passes, then the `json-repair` library as a last resort). The ladder is cheap in the common case — most posters succeed on step 1 and never pay for the retries.
 
+**The `hit_eos` trigger.** `_generate()` returns `(text, hit_eos)`; `hit_eos` is False when generation reached `max_new_tokens` without the brace processor ever letting an EOS through — the signature of a runaway that never closed the JSON. This condition is required in addition to the parse-error check because `_robust_json_parse()` is strong enough to salvage a truncated runaway into a *parseable but gutted* object (e.g. a single section), which carries no `"error"` key and would otherwise skip the retry that actually repairs it.
+
+**Why a repetition penalty only on retry.** A few posters drive the stock 8B into a repetition loop — it re-emits the `creators` array indefinitely and never closes the top-level object, grinding to the 18k cap and yielding `[ERR]`. A gentle `repetition_penalty` breaks the loop. It is applied **only on the retry/fallback passes**: healthy posters emit EOS on the primary pass and never reach the retry, so their output stays byte-identical to plain greedy. The value **1.15** was tuned empirically — 1.3 overshoots, corrupting the verbatim transcript the task depends on (`Genta -> Gentaa`, `drug polymer interactions -> drugpolymerinteractions`), while 1.15 breaks the loop and recovers a clean, fully-sectioned extraction. On the poster that motivated this (16083265), it moved the end-to-end result from `[ERR]` (0 sections) to a passing extraction (16/16 sections, rougeL 0.86).
+
 ## Quantization
 
 Default is **4-bit NF4**:
@@ -96,7 +101,7 @@ Flash Attention 2 is used automatically when `flash_attn` is importable; otherwi
 ## What we deliberately did *not* do
 
 - **No fine-tuning.** See note at top — stock instruct weights.
-- **No sampling.** Greedy only; reproducibility beats diversity for this task.
+- **No sampling.** Greedy only (`do_sample=False`); reproducibility beats diversity for this task. Generation stays deterministic even on the retry: `repetition_penalty` reshapes the greedy logits but adds no randomness, so the same poster still yields byte-identical JSON every run.
 - **No `min_new_tokens` for completeness.** It only guards one EOS token; the brace processor is the correct mechanism.
 - **No silent input truncation.** Over-long posters error out instead of being cut.
 

poster2json/extract.py

@@ -47,6 +47,18 @@
 MAX_RETRY_TOKENS = 24000
 MAX_INPUT_TOKENS = 15000
 
+# Runaway-loop guard. Some posters send the stock 8B into a repetition loop
+# (re-emitting the creators array) that never closes the JSON and grinds to the
+# token cap. The primary pass stays byte-identical (penalty 1.0); a runaway is
+# detected by the primary generation hitting the token cap without ever emitting
+# EOS, which triggers the retry. The retry/fallback passes apply a gentle
+# repetition penalty that breaks the loop while leaving the verbatim transcript
+# intact -- 1.15 was tuned empirically: 1.3 overshoots and corrupts author names
+# ("Genta" -> "Gentaa") and concatenates words, while 1.15 recovers a clean,
+# fully-sectioned extraction. A healthy poster emits EOS on the primary pass and
+# never reaches the retry, so its output is unchanged.
+RETRY_REPETITION_PENALTY = 1.15
+
 # Schema URL
 SCHEMA_URL = "https://posters.science/schema/v0.2/poster_schema.json"
 
@@ -1511,36 +1523,56 @@ def __call__(self, input_ids, scores):
         return scores
 
 
-def _generate(model, tokenizer, prompt: str, max_tokens: int) -> str:
-    """Generate response using the Llama model."""
+def _generate(
+    model, tokenizer, prompt: str, max_tokens: int, repetition_penalty: float = 1.0
+):
+    """Generate response using the Llama model.
+
+    repetition_penalty defaults to 1.0 (a no-op: identical to plain greedy) so
+    the primary pass is unchanged. The retry/fallback passes raise it to break
+    runaway repetition loops on posters whose primary pass never closed the JSON.
+
+    Returns (text, hit_eos). hit_eos is False when generation ran to the token
+    cap without emitting an EOS token -- the signature of a runaway that never
+    closed the JSON -- which the retry ladder uses to decide whether to retry.
+    """
     messages = [{"role": "user", "content": prompt}]
     input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
     input_length = inputs["input_ids"].shape[1]
 
-    log(f"Generating with max_tokens={max_tokens}, input_length={input_length}")
+    log(
+        f"Generating with max_tokens={max_tokens}, input_length={input_length}, "
+        f"repetition_penalty={repetition_penalty}"
+    )
 
     eos_ids = _get_eos_token_ids(tokenizer)
     processor = _JsonBraceProcessor(eos_ids, tokenizer, input_length)
 
-    streamer = ProgressStreamer(tokenizer, log_every=200)
+    gen_kwargs = dict(
+        max_new_tokens=max_tokens,
+        do_sample=False,
+        pad_token_id=tokenizer.eos_token_id,
+        streamer=ProgressStreamer(tokenizer, log_every=200),
+        logits_processor=LogitsProcessorList([processor]),
+    )
+    # Only pass the kwarg when it changes behaviour, so the primary pass stays
+    # byte-identical to the pre-fix output.
+    if repetition_penalty and repetition_penalty != 1.0:
+        gen_kwargs["repetition_penalty"] = repetition_penalty
+
     t0 = time.time()
     with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=max_tokens,
-            do_sample=False,
-            pad_token_id=tokenizer.eos_token_id,
-            streamer=streamer,
-            logits_processor=LogitsProcessorList([processor]),
-        )
+        outputs = model.generate(**inputs, **gen_kwargs)
     elapsed = time.time() - t0
     tokens_generated = outputs.shape[1] - input_length
+    hit_eos = outputs[0][-1].item() in eos_ids
     log(
-        f"   Generated {tokens_generated} tokens in {elapsed:.2f}s ({tokens_generated/elapsed:.1f} tok/s)"
+        f"   Generated {tokens_generated} tokens in {elapsed:.2f}s "
+        f"({tokens_generated/elapsed:.1f} tok/s), hit_eos={hit_eos}"
     )
 
-    return tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True)
+    return tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True), hit_eos
 
 
 # ============================
@@ -3137,24 +3169,35 @@ def extract_json_with_retry(
     prompt = EXTRACTION_PROMPT.format(raw_text=raw_text)
 
     log("Starting primary JSON extraction with full prompt")
-    response = _generate(model, tokenizer, prompt, MAX_JSON_TOKENS)
+    response, hit_eos = _generate(model, tokenizer, prompt, MAX_JSON_TOKENS)
     result = _as_result_dict(_robust_json_parse(response))
     if "error" in result:
         log(f"Primary JSON parse error: {result['error']}")
     else:
         log("Primary JSON parse succeeded")
 
-    # Retry with more tokens if truncation detected
-    if "error" in result or _is_truncated(result.get("raw", "")):
+    # Retry if the primary pass failed to parse, looks truncated, or ran to the
+    # token cap without emitting EOS (a runaway loop). The last condition is
+    # essential: the robust parser can salvage a truncated runaway into a
+    # parseable-but-gutted object (e.g. a single section), which has no "error"
+    # key and would otherwise skip the retry that actually fixes it. Apply the
+    # repetition penalty on the retry (not the primary pass) to break the loop.
+    if "error" in result or not hit_eos or _is_truncated(result.get("raw", "")):
         log(f"Retrying with max_tokens={MAX_RETRY_TOKENS}")
-        response = _generate(model, tokenizer, prompt, MAX_RETRY_TOKENS)
+        response, hit_eos = _generate(
+            model, tokenizer, prompt, MAX_RETRY_TOKENS,
+            repetition_penalty=RETRY_REPETITION_PENALTY,
+        )
         result = _as_result_dict(_robust_json_parse(response))
 
     # Fallback to shorter prompt
-    if "error" in result or _is_truncated(result.get("raw", "")):
+    if "error" in result or not hit_eos or _is_truncated(result.get("raw", "")):
         log("Using fallback shorter prompt")
         fallback_prompt = FALLBACK_PROMPT.format(raw_text=raw_text)
-        response = _generate(model, tokenizer, fallback_prompt, MAX_RETRY_TOKENS)
+        response, hit_eos = _generate(
+            model, tokenizer, fallback_prompt, MAX_RETRY_TOKENS,
+            repetition_penalty=RETRY_REPETITION_PENALTY,
+        )
         result = _as_result_dict(_robust_json_parse(response))
 
     result = _postprocess_json(