examples/loop_guard.py

"""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:     # empty rows, repeated header cells: tables repeat legitimately
                span, need = max(span, self.table_span), self.table_reps
            if r >= need and r * p >= span:
                return n - r * p + p          # keep one copy of the unit
        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]                     # the last line may still be growing
        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:])