| """Stop a page read that has started to loop, at serve time. |
| |
| A looping page runs to max_tokens (4,096): on the proxy models those pages take a third or more of all |
| decode time, and the tokens after the loop starts are garbage that only raise CER. The guard watches |
| the text as it streams and reports where the loop began; the client then aborts the request (vLLM |
| frees its slot) and keeps the text up to the end of the first copy. |
| |
| Two signals, both on the tail of the text: |
| period the tail is one unit of p characters repeated r times (exact), with p*r over a length that |
| grows for short units: dot leaders and "-----" rules are legitimate short repeats. Units with |
| table markup need far more (12 copies, 1,500 chars): empty rows and repeated header cells |
| are real table content. |
| lines the same line (stripped, at least `min_line` chars) repeated `line_reps` times in a row. |
| Tables and lists repeat structure, not whole identical lines. |
| Thresholds were set on eval v2 outputs (`python -m lightonocr_test.loop_guard tune ...`): no cut on a |
| page that stopped on its own (v1-ipo2), while catching most pages that hit the cap. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import sys |
|
|
|
|
| class LoopGuard: |
| def __init__(self, min_span: int = 240, short_unit: int = 12, short_span: int = 600, |
| table_span: int = 1500, table_reps: int = 12, |
| line_reps: int = 4, min_line: int = 12, every: int = 32): |
| self.min_span, self.short_unit, self.short_span = min_span, short_unit, short_span |
| self.table_span, self.table_reps = table_span, table_reps |
| self.line_reps, self.min_line, self.every = line_reps, min_line, every |
| self._checked = 0 |
|
|
| def _period(self, t: str) -> int | None: |
| n = len(t) |
| for p in range(1, 401): |
| span = self.short_span if p <= self.short_unit else self.min_span |
| if 3 * p > n or span > n: |
| continue |
| unit = t[n - p:] |
| r = 1 |
| while (r + 1) * p <= n and t[n - (r + 1) * p:n - r * p] == unit: |
| r += 1 |
| need = 3 |
| if "<t" in unit or "</t" in unit: |
| span, need = max(span, self.table_span), self.table_reps |
| if r >= need and r * p >= span: |
| return n - r * p + p |
| return None |
|
|
| def _lines(self, t: str) -> int | None: |
| lines = t.split("\n") |
| if len(lines) < self.line_reps + 1: |
| return None |
| body = lines[:-1] |
| last = body[-1].strip() |
| if len(last) < self.min_line: |
| return None |
| r = 1 |
| while r < len(body) and body[-1 - r].strip() == last: |
| r += 1 |
| if r >= self.line_reps: |
| keep = "\n".join(lines[:len(body) - r + 1]) |
| return len(keep) |
| return None |
|
|
| def check(self, text: str) -> int | None: |
| """Cut position if `text` ends in a loop, else None. Cheap to call on every streamed chunk: |
| it only looks again after `every` new characters.""" |
| if len(text) - self._checked < self.every: |
| return None |
| self._checked = len(text) |
| tail_start = max(0, len(text) - 4000) |
| tail = text[tail_start:] |
| for f in (self._period, self._lines): |
| c = f(tail) |
| if c is not None: |
| return tail_start + c |
| return None |
|
|
|
|
| def scan(text: str, step: int = 16, **kw) -> int | None: |
| """Replay a finished output in `step`-character chunks, as a stream would arrive.""" |
| g = LoopGuard(**kw) |
| for end in range(step, len(text) + step, step): |
| c = g.check(text[:end]) |
| if c is not None: |
| return c |
| return None |
|
|
|
|
| def tune(paths: list[str]) -> None: |
| """Per eval file: cuts on pages that stopped on their own (false alarms) and on capped pages, |
| characters saved, and CER before/after cutting (needs the references: eval_by_length sets).""" |
| from rapidfuzz.distance import Levenshtein |
|
|
| from .eval_by_length import sets, texts |
| refs = {} |
| for name, rows in sets(["v2"]).items(): |
| for k, (_, ref, meta) in enumerate(rows): |
| refs[(name, k)] = (ref, meta) |
| for p in paths: |
| rows = [json.loads(l) for l in open(p, encoding="utf-8")] |
| fa = caught = capped = 0 |
| chars = saved = 0 |
| cer0 = cer1 = 0.0 |
| n = 0 |
| for r in rows: |
| h = r.get("hyp") or "" |
| c = scan(h) |
| chars += len(h) |
| if r["capped"]: |
| capped += 1 |
| caught += c is not None |
| elif c is not None: |
| fa += 1 |
| key = (r["set"], r["k"]) |
| if key in refs: |
| ref, meta = refs[key] |
| R, H0 = texts(ref, meta, h) |
| _, H1 = texts(ref, meta, h[:c] if c is not None else h) |
| cer0 += Levenshtein.normalized_distance(R, H0) |
| cer1 += Levenshtein.normalized_distance(R, H1) |
| n += 1 |
| if c is not None: |
| saved += len(h) - c |
| print(f"{p}: pages {len(rows)} | false cuts {fa} of {len(rows) - capped} stopped pages | " |
| f"caught {caught} of {capped} capped | chars saved {100 * saved / max(1, chars):.1f}% | " |
| f"CER {cer0 / max(1, n):.4f} -> {cer1 / max(1, n):.4f}") |
|
|
|
|
| if __name__ == "__main__": |
| if sys.argv[1] == "tune": |
| tune(sys.argv[2:]) |
|
|
