diff --git a/README.md b/README.md index 193f755..35c0663 100644 --- a/README.md +++ b/README.md @@ -17,8 +17,9 @@ It keeps olmOCR's document management, queueing, and CLI 1:1, with two additions SOTA 1B Markdown-OCR model — image-only prompt, rendered at 1540px by default; `lightonocr2-soup` is the same adapter on the more-robust [`-ocr-soup`](https://huggingface.co/lightonai/LightOnOCR-2-1B-ocr-soup) - merged checkpoint; `glm-ocr`, `qianfan-ocr`, `infinity-parser2-flash`, and - `surya2` add four more document-OCR VLMs, documented below). + merged checkpoint; `glm-ocr`, `qianfan-ocr`, `infinity-parser2-flash`, + `surya2`, `unlimited-ocr`, and `ovisocr2` add six more document-OCR VLMs, + documented below). - **Opt-out resume** — completed work items are skipped on restart by default (olmOCR's done-flag behavior). `--no-resume` wipes prior progress and reprocesses the workspace from scratch. @@ -69,10 +70,19 @@ poetry run paperscale ./workspace \ flags, `results/*.jsonl`, and (with `--markdown`) `markdown/` live. - `--pdfs` — local PDF/image paths, a glob (`'docs/*.pdf'`), `.tar.gz` tarballs, or a `.txt` file listing one path per line. -- `--ocr-model {glm-ocr,infinity-parser2-flash,lightonocr2,lightonocr2-soup,markdown,olmocr,qianfan-ocr,surya2}` +- `--ocr-model {glm-ocr,infinity-parser2-flash,lightonocr2,lightonocr2-soup,markdown,olmocr,ovisocr2,qianfan-ocr,surya2,unlimited-ocr}` — which OCR adapter drives prompting/parsing. - `--model` — the served model id sent in each request (or a Hugging Face path for the internal server). +- `--disable-quality-check CHECK` — turn off one deterministic quality check so + it can no longer reject a page. Repeatable; `all` disables the gate entirely. + Checks: `empty_output`, `mojibake`, `control_characters`, + `refusal_boilerplate`, `malformed_frontmatter`, `repeated_character`, + `repeated_ngram`, `repeated_tail`, `truncation_indicator`, `length_anomaly`. + Prefer this over raising `--max_page_error_rate` when a corpus trips one + specific gate: it keeps the other checks armed instead of accepting every + failure mode at once. Pair it with `PAPERSCALE_DEBUG_REJECTS=.jsonl` to + see what a gate was dropping before you switch it off. Inputs are added to the workspace queue, grouped into work items, and processed concurrently by `--workers`. Re-running the same command **resumes**: finished @@ -240,6 +250,35 @@ poetry run paperscale ./workspace --pdfs './docs/*.pdf' \ --ocr-model surya2 --server http://127.0.0.1:8000/v1 --markdown ``` +### OvisOCR2 + +`--ocr-model ovisocr2` drives +[OvisOCR2](https://huggingface.co/ATH-MaaS/OvisOCR2) (`ATH-MaaS/OvisOCR2`, 0.8B, +`Qwen3_5ForConditionalGeneration`), a compact end-to-end page parser post-trained +from Qwen3.5-0.8B. It scores 96.58 on OmniDocBench v1.6 — the first end-to-end +model to top that leaderboard — and 75.06 Avg3 on PureDocBench. One call returns +the full page as Markdown, with LaTeX formulas and HTML `` markup; +paperscale strips the `` placeholders the model emits for +charts and figures (their crops are never written) and flags those pages as +diagrams instead. + +The vendor pins `vllm==0.22.1` and selects the Triton GDN prefill backend — the +model interleaves Gated-DeltaNet `linear_attention` layers with full attention: + +```bash +vllm serve ATH-MaaS/OvisOCR2 --port 8000 \ + --gdn-prefill-backend triton \ + --limit-mm-per-prompt '{"image": 1}' \ + --mm-processor-kwargs '{"images_kwargs": {"min_pixels": 200704, "max_pixels": 8294400}}' + +poetry run paperscale ./workspace --pdfs './docs/*.pdf' \ + --ocr-model ovisocr2 --server http://127.0.0.1:8000/v1 --markdown +``` + +Pages render at 1540px by default, comfortably inside the vendor's +`448²`–`2880²` pixel band; that band tolerates roughly 3270px on the long edge +via `--target_longest_image_dim` if a corpus has fine print. + ## Outputs For each work item, paperscale writes `workspace/results/output_.jsonl` — diff --git a/src/paperscale/models/__init__.py b/src/paperscale/models/__init__.py index 1002e0f..c86f352 100644 --- a/src/paperscale/models/__init__.py +++ b/src/paperscale/models/__init__.py @@ -12,6 +12,7 @@ from paperscale.models.lightonocr import LightOnOCRModel, LightOnOCRSoupModel from paperscale.models.markdown import MarkdownModel from paperscale.models.olmocr import OlmOCRModel +from paperscale.models.ovisocr2 import OvisOCR2Model from paperscale.models.qianfan import QianfanOCRModel from paperscale.models.surya import Surya2Model from paperscale.models.unlimited_ocr import UnlimitedOCRModel @@ -28,6 +29,7 @@ "infinity-parser2-flash": InfinityParser2FlashModel, "surya2": Surya2Model, "unlimited-ocr": UnlimitedOCRModel, + "ovisocr2": OvisOCR2Model, } @@ -52,6 +54,7 @@ def build_ocr_model(name: str) -> OCRModel: "InfinityParser2FlashModel", "Surya2Model", "UnlimitedOCRModel", + "OvisOCR2Model", "MODEL_REGISTRY", "DEFAULT_MODEL", "build_ocr_model", diff --git a/src/paperscale/models/markdown.py b/src/paperscale/models/markdown.py index 4446e07..544faf2 100644 --- a/src/paperscale/models/markdown.py +++ b/src/paperscale/models/markdown.py @@ -17,6 +17,12 @@ _CODE_FENCE = re.compile(r"^\s*```(?:markdown|md)?\s*\n(?P.*?)\n```\s*$", re.DOTALL | re.IGNORECASE) +# A complete ... reasoning block. Several OCR VLMs emit one when +# thinking is enabled; paperscale never enables it, so stripping is defensive. +# Matched-pair only, so a literal "" transcribed from the page cannot +# truncate real content. +_THINK_BLOCK = re.compile(r".*?", re.DOTALL) + def _strip_code_fence(content: str) -> str: """Remove a single wrapping ```markdown ... ``` fence some models add.""" @@ -24,6 +30,16 @@ def _strip_code_fence(content: str) -> str: return match.group("body") if match else content +def strip_think_blocks(content: str) -> str: + """Drop any complete ``...`` reasoning block from a response.""" + return _THINK_BLOCK.sub("", content) + + +def has_html_table(markdown: str | None) -> bool: + """Whether Markdown carries an HTML ``
``, which paperscale keeps as-is.""" + return "`` for tables (paperscale accepts these as-is, like ``qianfan-ocr``), +and visual regions as placeholder ```` +tags with coordinates scaled to [0, 1000). paperscale never writes the crop +files those tags reference, so :meth:`parse` strips them and reports the page as +a diagram instead. + +vLLM serving (vendor pins ``vllm==0.22.1``; the checkpoint declares +``transformers_version 4.57.0.dev0``, so pair it with transformers >= 4.57. +Verified end-to-end on vLLM 0.27.1, which registers the arch natively):: + + vllm serve ATH-MaaS/OvisOCR2 --port 8000 \ + --gdn-prefill-backend triton \ + --limit-mm-per-prompt '{"image": 1}' \ + --mm-processor-kwargs '{"images_kwargs": {"min_pixels": 200704, "max_pixels": 8294400}}' + +``--gdn-prefill-backend triton`` is the vendor's ``gdn_prefill_backend="triton"`` +LLM kwarg; the mm-processor bounds are their ``min_pixels=448*448`` / +``max_pixels=2880*2880``. The default 1540px render sits comfortably inside that +band (~1.8 MP for a letter page), and the band tolerates up to ~3270px on the +long edge via ``--target_longest_image_dim`` if fine print needs it. + +Degenerate output: the vendor's own parser trims a repeating tail out of long +responses before returning them. This adapter deliberately does **not** — a +looping page is an *incomplete* page (the model spent its budget repeating +instead of transcribing the rest), and rewriting it here would hand the pipeline +something that reads as clean and complete. The same detection lives in the +quality gate instead, as +:func:`paperscale.quality.verifier._has_repeated_tail_loop`, so the page is +rejected and retried at a higher temperature. Disable it with +``--disable-quality-check repeated_tail`` if a corpus trips it systematically. +""" + +from __future__ import annotations + +import re +from dataclasses import replace + +from paperscale.models.markdown import MarkdownModel, has_html_table, strip_think_blocks +from paperscale.prompts import PageResponse + +# Vendor inference prompt, verbatim from the model card's OvisOCR2Parser. The +# leading newline and the literal ``{left}``/``{top}``/… braces are part of the +# string (the vendor builds it as a plain, non-f string). Keep both: the chat +# template trims only the *ends* of the rendered user turn, so with the image +# part sent first this newline survives as the image/instruction separator, +# reproducing the vendor's exact token sequence. +OVIS_OCR2_PROMPT = ( + "\nExtract all readable content from the image in natural human reading order " + "and output the result as a single Markdown document. For charts or images, " + 'represent them using an HTML image tag: , ' + "where left, top, right, bottom are bounding box coordinates scaled to [0, 1000). " + "Format formulas as LaTeX. Format tables as HTML:
...
. Transcribe " + "all other text as standard Markdown. Preserve the original text without " + "translation or paraphrasing." +) + +# Placeholder tags for visual regions: . +# The vendor's own parser drops these by default (filter_imgtags=True) because the +# referenced crops only exist if you save them alongside the Markdown; paperscale +# never does, so every occurrence — block-level or inline — is a dead reference. +_BBOX_IMG = re.compile(r'') + +# Blank lines left behind where a stripped tag was the whole block. +_EXCESS_BLANKS = re.compile(r"\n{3,}") + + +class OvisOCR2Model(MarkdownModel): + """Drives OvisOCR2, which transcribes a page image to full-page Markdown.""" + + default_model_name = "ATH-MaaS/OvisOCR2" + preferred_longest_image_dim = 1540 + + def __init__(self) -> None: + super().__init__(prompt=OVIS_OCR2_PROMPT) + + def build_messages(self, image_base64: str) -> list[dict]: + # Image part first, then the instruction: this is the vendor's content + # order, and it is what keeps OVIS_OCR2_PROMPT's leading newline from + # being trimmed off the front of the rendered user turn. + return [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}}, + {"type": "text", "text": self._prompt}, + ], + } + ] + + def sampling_params(self) -> dict: + # Vendor decodes with max_tokens=16384; full pages of dense HTML tables + # overrun the pipeline's 8000 default. Temperature stays pipeline-owned. + return {"max_tokens": 16384} + + def parse(self, content: str) -> PageResponse: + # Defensive: the chat template pre-closes thinking unless enable_thinking is + # explicitly true, which paperscale never sets, so no block should appear. + content = strip_think_blocks(content) + # A bbox img tag is the model's marker for a chart/figure region, so read + # the diagram flag off the raw text before the dead references go away. + is_diagram = bool(_BBOX_IMG.search(content)) + content = _EXCESS_BLANKS.sub("\n\n", _BBOX_IMG.sub("", content)) + result = super().parse(content) + # Tables come back as HTML, which paperscale keeps as-is; recompute the + # flag from the emitted Markdown rather than converting it. + return replace(result, is_table=has_html_table(result.natural_text), is_diagram=is_diagram) diff --git a/src/paperscale/models/qianfan.py b/src/paperscale/models/qianfan.py index bc89170..d8463c3 100644 --- a/src/paperscale/models/qianfan.py +++ b/src/paperscale/models/qianfan.py @@ -15,17 +15,11 @@ from __future__ import annotations -import re from dataclasses import replace -from paperscale.models.markdown import MarkdownModel +from paperscale.models.markdown import MarkdownModel, has_html_table, strip_think_blocks from paperscale.prompts import PageResponse -# A complete ... reasoning block (only emitted when -# enable_thinking=True, which paperscale never sets). Matched-pair only, so a -# literal "" transcribed from the page can't truncate real content. -_THINK_BLOCK = re.compile(r".*?", re.DOTALL) - class QianfanOCRModel(MarkdownModel): """Drives Qianfan-OCR, which transcribes a page image to reading-ordered Markdown.""" @@ -45,9 +39,8 @@ def parse(self, content: str) -> PageResponse: # Defensive: drop any ... reasoning block. paperscale never # sets enable_thinking, so this is normally a no-op; matched-pair only, so a # literal "" in the page can't truncate the transcription. - content = _THINK_BLOCK.sub("", content) + content = strip_think_blocks(content) result = super().parse(content) # Qianfan returns tables as HTML, which paperscale accepts as-is. Recompute # is_table from the emitted Markdown rather than convert it. - is_table = " DeterministicQualityVerifier: + """Build the page verifier, honouring ``--disable-quality-check``.""" + codes = expand_disabled_checks(disabled) + if codes: + logger.warning(f"Quality checks disabled: {', '.join(sorted(codes))}. Pages these would reject are now accepted as-is.") + return DeterministicQualityVerifier(disabled_checks=codes) + + # A page whose rendered image has less ink than this is treated as genuinely blank: # an empty/degenerate OCR result for it is accepted as a successful empty page # rather than retried. Blank scans sit ~0.002-0.005; content pages are higher. @@ -117,7 +128,7 @@ # Quality diagnostics that, on a near-blank render, mean "genuinely blank page" # (empty output, or a model repetition loop on noise) rather than a read failure. -_BLANK_ELIGIBLE_DIAGNOSTICS = frozenset({"empty_output", "repeated_ngram", "repeated_character"}) +_BLANK_ELIGIBLE_DIAGNOSTICS = frozenset({"empty_output", "repeated_ngram", "repeated_character", "repeated_tail"}) def _render_is_blank(image_base64: str) -> bool: @@ -235,7 +246,7 @@ async def try_single_page( model_response_markdown = base_response_data["choices"][0]["message"]["content"] page_response = args.ocr_model.parse(model_response_markdown) - finding = _verifier.classify(page_response.natural_text or "") + finding = args.quality_verifier.classify(page_response.natural_text or "") if finding.accepted: is_valid, is_terminal = True, False elif finding.kind in _BLANK_ELIGIBLE_DIAGNOSTICS and render_is_blank: @@ -1167,6 +1178,20 @@ def _build_arg_parser() -> argparse.ArgumentParser: parser.add_argument("--pages_per_group", type=int, default=100, help="Aim for this many PDF pages per work item group.") parser.add_argument("--max_page_retries", type=int, default=8, help="Max number of times to retry a page.") parser.add_argument("--max_page_error_rate", type=float, default=0.004, help="Allowable fraction of fallback pages per document.") + parser.add_argument( + "--disable-quality-check", + dest="disabled_quality_checks", + action="append", + default=[], + metavar="CHECK", + choices=(*QUALITY_CHECK_CODES, DISABLE_ALL_CHECKS), + help="Skip a deterministic quality check so it cannot reject a page. Repeatable. " + f"'{DISABLE_ALL_CHECKS}' disables the gate entirely. Choices: {', '.join(QUALITY_CHECK_CODES)}.", + ) + # main() replaces this from --disable-quality-check; the default keeps the fully + # armed gate available to callers that drive try_single_page without going + # through main() (tests, embedders). + parser.set_defaults(quality_verifier=_verifier) parser.add_argument("--workers", type=int, default=4, help="Max number of page groups processed at once.") parser.add_argument("--max_concurrent_requests", type=int, default=500, help="Max requests in-flight to the inference provider at once.") parser.add_argument("--max_server_ready_timeout", type=int, default=600, help="Seconds to wait for the server to become ready.") @@ -1244,6 +1269,7 @@ async def main(): # Resolve the OCR model adapter and the served model name. args.ocr_model = build_ocr_model(args.ocr_model_name) + args.quality_verifier = _build_verifier(args.disabled_quality_checks) if args.model is None: args.model = args.ocr_model.default_model_name # Render size: explicit --target_longest_image_dim wins, else the model's preferred. diff --git a/src/paperscale/quality/verifier.py b/src/paperscale/quality/verifier.py index 775f4d2..cfc3f1b 100644 --- a/src/paperscale/quality/verifier.py +++ b/src/paperscale/quality/verifier.py @@ -4,6 +4,7 @@ from collections import Counter from dataclasses import dataclass, field +import math import re _CONTROL_CHAR_RE = re.compile(r"[\x00-\x08\x0b\x0c\x0e-\x1f]") @@ -57,14 +58,41 @@ class VerificationFinding: warnings: list[str] = field(default_factory=list) +# Every check ``assess_markdown_fragment`` can raise, in the order it runs them. +# This is the vocabulary accepted by ``--disable-quality-check``. +QUALITY_CHECK_CODES = ( + "empty_output", + "mojibake", + "control_characters", + "refusal_boilerplate", + "malformed_frontmatter", + "repeated_character", + "repeated_ngram", + "repeated_tail", + "truncation_indicator", + "length_anomaly", +) + +#: Accepted by ``--disable-quality-check`` alongside the individual codes. +DISABLE_ALL_CHECKS = "all" + + +def expand_disabled_checks(names: list[str]) -> frozenset[str]: + """Turn ``--disable-quality-check`` values into the set of codes to skip.""" + if DISABLE_ALL_CHECKS in names: + return frozenset(QUALITY_CHECK_CODES) + return frozenset(names) + + class DeterministicQualityVerifier: """Local, no-extra-model verifier for v1 page Markdown fragments.""" - def __init__(self, optional_slm: object | None = None) -> None: + def __init__(self, optional_slm: object | None = None, disabled_checks: frozenset[str] = frozenset()) -> None: self.optional_slm = optional_slm + self.disabled_checks = disabled_checks def classify(self, markdown: str) -> VerificationFinding: - report = assess_markdown_fragment(markdown) + report = assess_markdown_fragment(markdown, disabled_checks=self.disabled_checks) if report.accepted: return VerificationFinding(True, "ok", "none", warnings=[]) issue = report.issues[0] @@ -73,39 +101,52 @@ def classify(self, markdown: str) -> VerificationFinding: return VerificationFinding(False, kind, retry_class, warnings=[]) -def assess_markdown_fragment(markdown: str) -> QualityReport: - """Assess whether a page Markdown fragment is coherent enough to assemble.""" +def assess_markdown_fragment(markdown: str, disabled_checks: frozenset[str] = frozenset()) -> QualityReport: + """Assess whether a page Markdown fragment is coherent enough to assemble. + + ``disabled_checks`` names codes from :data:`QUALITY_CHECK_CODES` to skip, so a + corpus that trips a gate systematically can be processed without it rather + than by loosening ``--max_page_error_rate`` for every failure mode at once. + """ issues: list[QualityIssue] = [] text = markdown.strip() + def enabled(code: str) -> bool: + return code not in disabled_checks + if not text: - issues.append(QualityIssue("empty_output", "OCR output is empty or whitespace only.")) - return QualityReport(accepted=False, severity="error", issues=issues) + if enabled("empty_output"): + issues.append(QualityIssue("empty_output", "OCR output is empty or whitespace only.")) + return QualityReport(accepted=False, severity="error", issues=issues) + return QualityReport(accepted=True, severity="ok", issues=issues) replacement_count = text.count("\ufffd") - if replacement_count >= 3 or replacement_count / max(len(text), 1) > 0.02: + if enabled("mojibake") and (replacement_count >= 3 or replacement_count / max(len(text), 1) > 0.02): issues.append(QualityIssue("mojibake", "OCR output contains too many Unicode replacement characters.")) - if _CONTROL_CHAR_RE.search(text): + if enabled("control_characters") and _CONTROL_CHAR_RE.search(text): issues.append(QualityIssue("control_characters", "OCR output contains control characters.")) - if _has_refusal_boilerplate(text): + if enabled("refusal_boilerplate") and _has_refusal_boilerplate(text): issues.append(QualityIssue("refusal_boilerplate", "OCR output contains refusal boilerplate.")) - if _has_malformed_frontmatter(text): + if enabled("malformed_frontmatter") and _has_malformed_frontmatter(text): issues.append(QualityIssue("malformed_frontmatter", "OCR output has malformed frontmatter or schema preamble.")) - if _has_repeated_character_run(text): + if enabled("repeated_character") and _has_repeated_character_run(text): issues.append(QualityIssue("repeated_character", "OCR output contains an abnormal repeated-character run.")) - if _has_repeated_ngram_loop(text): + if enabled("repeated_ngram") and _has_repeated_ngram_loop(text): issues.append(QualityIssue("repeated_ngram", "OCR output appears to repeat the same phrase.")) - if _has_truncation_indicator(text): + if enabled("repeated_tail") and _has_repeated_tail_loop(text): + issues.append(QualityIssue("repeated_tail", "OCR output ends in a repeating loop.")) + + if enabled("truncation_indicator") and _has_truncation_indicator(text): issues.append(QualityIssue("truncation_indicator", "OCR output appears truncated.")) - if _has_length_anomaly(text): + if enabled("length_anomaly") and _has_length_anomaly(text): issues.append(QualityIssue("length_anomaly", "OCR output length looks anomalous.")) accepted = not any(issue.severity == "error" for issue in issues) @@ -161,6 +202,58 @@ def _has_repeated_ngram_loop(text: str) -> bool: return False +# Degeneration that runs to the end of the output: the model got stuck emitting +# one unit until it ran out of budget. _has_repeated_ngram_loop cannot see this +# once the unit grows: it scores ``count * ngram_size / len(tokens)`` with +# ngram_size capped at 5, so the score tops out at ``5 / period`` and a looped +# *sentence* (13+ tokens) stays under the 0.35 threshold no matter how much of +# the page it eats — measured at 0.35 even when the loop is 90% of the output. +# Walking back from the end at a fixed period instead makes period length +# irrelevant. The period/repeat/length constants are OvisOCR2's vendor cleaner +# (https://huggingface.co/ATH-MaaS/OvisOCR2), reused here as a *detector* so the +# page is retried rather than silently rewritten. +_TAIL_LOOP_MAX_PERIOD = 200 +_TAIL_LOOP_MIN_REPEATS = 5 +_TAIL_LOOP_MIN_CHARS = 100 +# The loop must also dominate the output, reusing the n-gram gate's 0.35 share. +# This is what separates a genuine loop from legitimately tiling content — a form +# whose trailing table rows are identical blanks sits near 0.23 and is kept. +_TAIL_LOOP_MIN_SHARE = 0.35 +# The loop need not reach the final character. A model that loops table rows and +# then closes the tag leaves "" behind it, and OvisOCR2 emits HTML tables, +# so that is the common shape — anchoring at len(text) misses it entirely. Probe a +# ladder of end positions so a trailer of ordinary text after the loop cannot hide +# it. Stepping back one character at a time does not work: on prose a character +# coincidentally lines up at the period within a few steps, which stops the search +# inside the trailer rather than at the end of the loop. +_TAIL_LOOP_TRAILERS = (0, 8, 24, 64, 160, 400) + + +def _has_repeated_tail_loop(text: str) -> bool: + length = len(text) + if length < _TAIL_LOOP_MIN_CHARS: + return False + # All three thresholds grow monotonically with the length of the repeating run, + # so for a given period and end position only ``span`` characters decide the + # answer: if that slice repeats at ``period``, every threshold is met, and if it + # does not, no longer run ends there either. Testing it as one slice equality + # keeps the scan in C — walking character by character in Python costs ~300ms on + # a page whose tail repeats just under the share threshold. + share_floor = math.ceil(_TAIL_LOOP_MIN_SHARE * length) + for trailer in _TAIL_LOOP_TRAILERS: + end = length - trailer + if end < _TAIL_LOOP_MIN_CHARS: + break # trailers only grow, so no later one leaves room either + for period in range(1, _TAIL_LOOP_MAX_PERIOD + 1): + span = max(_TAIL_LOOP_MIN_CHARS, _TAIL_LOOP_MIN_REPEATS * period, share_floor) + if span > end: + break # span only grows with period, so no larger period fits either + tail = text[end - span : end] + if tail[period:] == tail[:-period]: + return True + return False + + def _has_refusal_boilerplate(text: str) -> bool: lowered = text.lower() if any(re.search(pattern, lowered) for pattern in _REFUSAL_PATTERNS): diff --git a/tests/test_models.py b/tests/test_models.py index 423c6eb..642610f 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -11,6 +11,7 @@ LightOnOCRSoupModel, MarkdownModel, OlmOCRModel, + OvisOCR2Model, QianfanOCRModel, Surya2Model, UnlimitedOCRModel, @@ -18,6 +19,7 @@ ) from paperscale.models.glmocr import GLM_OCR_PROMPT from paperscale.models.markdown import _strip_code_fence +from paperscale.models.ovisocr2 import OVIS_OCR2_PROMPT from paperscale.prompts import PageResponse @@ -35,6 +37,7 @@ def test_build_known_models(self): self.assertIsInstance(build_ocr_model("infinity-parser2-flash"), InfinityParser2FlashModel) self.assertIsInstance(build_ocr_model("surya2"), Surya2Model) self.assertIsInstance(build_ocr_model("unlimited-ocr"), UnlimitedOCRModel) + self.assertIsInstance(build_ocr_model("ovisocr2"), OvisOCR2Model) def test_registry_contents(self): self.assertEqual( @@ -49,6 +52,7 @@ def test_registry_contents(self): "infinity-parser2-flash", "surya2", "unlimited-ocr", + "ovisocr2", }, ) @@ -120,31 +124,14 @@ def test_has_guided_regex(self): self.assertIsNotNone(self.model.guided_regex()) def test_parse_front_matter(self): - content = ( - "---\n" - "primary_language: en\n" - "is_rotation_valid: True\n" - "rotation_correction: 0\n" - "is_table: False\n" - "is_diagram: False\n" - "---\n" - "Hello world" - ) + content = "---\nprimary_language: en\nis_rotation_valid: True\nrotation_correction: 0\nis_table: False\nis_diagram: False\n---\nHello world" result = self.model.parse(content) self.assertEqual(result.natural_text, "Hello world") self.assertEqual(result.primary_language, "en") self.assertTrue(result.is_rotation_valid) def test_parse_rotation_flag(self): - content = ( - "---\n" - "primary_language: null\n" - "is_rotation_valid: False\n" - "rotation_correction: 90\n" - "is_table: False\n" - "is_diagram: False\n" - "---\n" - ) + content = "---\nprimary_language: null\nis_rotation_valid: False\nrotation_correction: 90\nis_table: False\nis_diagram: False\n---\n" result = self.model.parse(content) self.assertFalse(result.is_rotation_valid) self.assertEqual(result.rotation_correction, 90) @@ -421,7 +408,7 @@ def test_parse_flags_tolerate_attribute_quoting(self): # is_table/is_diagram must not depend on the model's quoting style. single = "
x
" self.assertTrue(self.model.parse(single).is_table) - spaced = "
fig
" + spaced = '
fig
' self.assertTrue(self.model.parse(spaced).is_diagram) @@ -449,10 +436,7 @@ def test_build_messages_has_image_prefixed_prompt(self): self.assertEqual(image_part["image_url"]["url"], "data:image/png;base64,QUJD") def test_parse_unwraps_ref_and_drops_det(self): - raw = ( - "<|ref|># Heading<|/ref|><|det|>[[10,20,900,60]]<|/det|>\n\n" - "Body text with <|ref|>a grounded span<|/ref|><|det|>[[1,2,3,4]]<|/det|> inline." - ) + raw = "<|ref|># Heading<|/ref|><|det|>[[10,20,900,60]]<|/det|>\n\nBody text with <|ref|>a grounded span<|/ref|><|det|>[[1,2,3,4]]<|/det|> inline." result = self.model.parse(raw) self.assertEqual( result.natural_text, @@ -474,5 +458,79 @@ def test_parse_empty_page_is_none(self): self.assertIsNone(self.model.parse("<|ref|><|/ref|><|det|>[[0,0,0,0]]<|/det|>").natural_text) +class OvisOCR2ModelTests(unittest.TestCase): + def setUp(self): + self.model = OvisOCR2Model() + + def test_recipe(self): + self.assertEqual(self.model.default_model_name, "ATH-MaaS/OvisOCR2") + self.assertEqual(self.model.preferred_longest_image_dim, 1540) + # Vendor decodes at 16384 tokens; temperature stays pipeline-owned. + self.assertEqual(self.model.sampling_params(), {"max_tokens": 16384}) + self.assertNotIn("temperature", self.model.sampling_params()) + self.assertIsNone(self.model.guided_regex()) + + def test_prompt_is_vendor_verbatim(self): + # The leading newline separates the image from the instruction once the + # chat template renders the turn; the braces are literal, not a format spec. + self.assertTrue(OVIS_OCR2_PROMPT.startswith("\nExtract all readable content from the image")) + self.assertIn('', OVIS_OCR2_PROMPT) + self.assertIn("Format tables as HTML: ...
.", OVIS_OCR2_PROMPT) + self.assertTrue(OVIS_OCR2_PROMPT.endswith("without translation or paraphrasing.")) + + def test_build_messages_puts_image_before_prompt(self): + # Order is load-bearing: the chat template trims the ends of the rendered + # user turn, so a text-first layout would eat the prompt's leading newline. + messages = self.model.build_messages("QUJD") + self.assertEqual(len(messages), 1) + image_part, text_part = messages[0]["content"] + self.assertEqual(image_part["image_url"]["url"], "data:image/png;base64,QUJD") + self.assertEqual(text_part, {"type": "text", "text": OVIS_OCR2_PROMPT}) + + def test_parse_passes_markdown_through(self): + result = self.model.parse("# Title\n\nBody text.") + self.assertEqual(result.natural_text, "# Title\n\nBody text.") + self.assertTrue(result.is_rotation_valid) + self.assertEqual(result.rotation_correction, 0) + self.assertFalse(result.is_table) + self.assertFalse(result.is_diagram) + + def test_parse_drops_bbox_image_tags_and_flags_diagram(self): + raw = '# Report\n\n\n\nFigure 1 shows the trend.' + result = self.model.parse(raw) + self.assertEqual(result.natural_text, "# Report\n\nFigure 1 shows the trend.") + self.assertNotIn(" below.') + self.assertEqual(result.natural_text, "See below.") + self.assertTrue(result.is_diagram) + + def test_parse_detects_html_table(self): + result = self.model.parse("
1
") + self.assertTrue(result.is_table) + self.assertFalse(result.is_diagram) + + def test_parse_strips_think_block(self): + result = self.model.parse("\n\n\n\n# Heading") + self.assertEqual(result.natural_text, "# Heading") + + def test_parse_keeps_content_before_stray_think_close(self): + # Matched-pair only: a transcribed "" must not truncate the page. + result = self.model.parse("Body mentioning literally.") + self.assertEqual(result.natural_text, "Body mentioning literally.") + + def test_parse_empty_page_is_none(self): + self.assertIsNone(self.model.parse("").natural_text) + # A page whose only output was a figure placeholder is empty text, but is + # still a diagram page. + figure_only = self.model.parse('') + self.assertIsNone(figure_only.natural_text) + self.assertTrue(figure_only.is_diagram) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_pipeline_units.py b/tests/test_pipeline_units.py index e51b72f..3969492 100644 --- a/tests/test_pipeline_units.py +++ b/tests/test_pipeline_units.py @@ -1,6 +1,7 @@ """Tests for pure pipeline helpers (no server, no rendering).""" import errno +import inspect import logging import os import tempfile @@ -11,6 +12,7 @@ from paperscale import pipeline from paperscale.pipeline import PageResult, _build_arg_parser, _tui_log_path, classify_document, count_documents, count_retries +from paperscale.quality.verifier import QUALITY_CHECK_CODES from paperscale.prompts import PageResponse from paperscale.tui import install_tui_logging, restore_console_logging @@ -651,5 +653,53 @@ async def boom(*a, **kw): self.assertIn(handler, log.handlers) +class DisableQualityCheckFlagTests(unittest.TestCase): + """--disable-quality-check, from argv through to the verifier a page is judged by.""" + + # A page dominated by one repeated 22-token sentence. The loop unit is long + # enough that repeated_ngram provably cannot score it (its ceiling is + # 5 / period), so repeated_tail is the *only* check that rejects this page — + # which is what makes it a clean probe for disabling that one check. + LOOPING_PAGE = ("The Court finds that the respondent failed to establish a prima facie case. " * 20) + ( + "Payment of the sum shall be made to the clerk within thirty (30) calendar days of the invoice date. " * 160 + ) + + def test_default_is_no_checks_disabled(self): + args = _build_arg_parser().parse_args(["/tmp/ws"]) + self.assertEqual(args.disabled_quality_checks, []) + # A run that never reaches main() still gets a fully armed gate. + self.assertFalse(args.quality_verifier.classify(self.LOOPING_PAGE).accepted) + + def test_flag_is_repeatable_and_reaches_the_verifier(self): + args = _build_arg_parser().parse_args(["/tmp/ws", "--disable-quality-check", "repeated_tail", "--disable-quality-check", "mojibake"]) + self.assertEqual(args.disabled_quality_checks, ["repeated_tail", "mojibake"]) + verifier = pipeline._build_verifier(args.disabled_quality_checks) + self.assertEqual(verifier.disabled_checks, frozenset({"repeated_tail", "mojibake"})) + finding = verifier.classify(self.LOOPING_PAGE) + self.assertTrue(finding.accepted, finding.kind) + self.assertNotEqual(finding.kind, "repeated_tail") + + def test_disabling_an_unrelated_check_leaves_the_page_rejected(self): + args = _build_arg_parser().parse_args(["/tmp/ws", "--disable-quality-check", "mojibake"]) + self.assertFalse(pipeline._build_verifier(args.disabled_quality_checks).classify(self.LOOPING_PAGE).accepted) + + def test_all_disables_every_check(self): + args = _build_arg_parser().parse_args(["/tmp/ws", "--disable-quality-check", "all"]) + verifier = pipeline._build_verifier(args.disabled_quality_checks) + self.assertEqual(verifier.disabled_checks, frozenset(QUALITY_CHECK_CODES)) + for text in ("", self.LOOPING_PAGE, "I cannot provide that."): + self.assertTrue(verifier.classify(text).accepted) + + def test_unknown_check_is_rejected_by_argparse(self): + with self.assertRaises(SystemExit): + _build_arg_parser().parse_args(["/tmp/ws", "--disable-quality-check", "not_a_check"]) + + def test_try_single_page_uses_the_verifier_on_args(self): + # The judged-by path: whatever verifier main() puts on args is the one + # try_single_page consults, with no silent fallback to the module default. + source = inspect.getsource(pipeline.try_single_page) + self.assertIn("args.quality_verifier.classify(", source) + + if __name__ == "__main__": unittest.main() diff --git a/tests/test_quality_verifier.py b/tests/test_quality_verifier.py index be91463..3fbac87 100644 --- a/tests/test_quality_verifier.py +++ b/tests/test_quality_verifier.py @@ -1,8 +1,17 @@ """Tests for the deterministic quality verifier that gates page acceptance.""" +import inspect +import re +import time import unittest -from paperscale.quality.verifier import DeterministicQualityVerifier, assess_markdown_fragment +from paperscale.quality.verifier import ( + QUALITY_CHECK_CODES, + DeterministicQualityVerifier, + _has_repeated_tail_loop, + assess_markdown_fragment, + expand_disabled_checks, +) class AssessMarkdownFragmentTests(unittest.TestCase): @@ -129,5 +138,146 @@ def test_mojibake_is_retryable(self): self.assertEqual(finding.retry_class, "retryable") +GOOD_SENTENCE = "The Court finds that the respondent failed to establish a prima facie case under the applicable provisions of the statute. " +# A sentence-length loop unit: 13 tokens, longer than the n-gram check can score. +LOOP_SENTENCE = "Payment shall be made within thirty (30) days of invoice. " + + +class RepeatedTailLoopTests(unittest.TestCase): + """The tail-loop check covers what the n-gram check structurally cannot. + + ``_has_repeated_ngram_loop`` scores ``count * ngram_size / len(tokens)`` with + ngram_size capped at 5, so it tops out at ``5 / period``: a looped sentence + never reaches the 0.35 threshold however much of the page it consumes. + """ + + def test_sentence_length_loop_rejected(self): + # 86% of this page is one repeated sentence; the n-gram check scores it + # at 0.35 and lets it through, so the tail check has to catch it. + text = (GOOD_SENTENCE * 20) + "\n\n" + (LOOP_SENTENCE * 260) + report = assess_markdown_fragment(text) + self.assertFalse(report.accepted) + self.assertTrue(any(issue.code == "repeated_tail" for issue in report.issues)) + + def test_loop_followed_by_a_trailer_is_still_rejected(self): + # The loop rarely runs to the final character. OvisOCR2 emits HTML tables, + # so a model looping table rows typically closes the tag afterwards; an + # end-anchored check would accept every one of these pages. + looping = (GOOD_SENTENCE * 20) + "\n\n" + (LOOP_SENTENCE * 260) + for trailer in ( + "", + "END OF DOCUMENT.", + " The auditor then reviewed the remaining ledger entries in detail.", + " The auditor reviewed further entries." * 10, + ): + with self.subTest(trailer=trailer[:24]): + report = assess_markdown_fragment(looping + trailer) + self.assertTrue(any(issue.code == "repeated_tail" for issue in report.issues)) + + def test_long_period_loop_rejected(self): + # A 22-token unit: the n-gram check cannot score above 5/22 = 0.23. + unit = "Payment of the sum shall be made to the clerk within thirty (30) calendar days of the invoice date. " + report = assess_markdown_fragment((GOOD_SENTENCE * 12) + "\n\n" + (unit * 120)) + self.assertFalse(report.accepted) + self.assertTrue(any(issue.code == "repeated_tail" for issue in report.issues)) + + def test_clean_page_accepted(self): + # Varied prose, not one sentence repeated: a page that really is the same + # sentence 80 times is degenerate, and the gate is right to reject it. + page = "# Memorandum of Decision\n\n" + "".join( + f"Paragraph {n} records that the witness described the events of the {n}th of March and produced exhibit {n} in support. " for n in range(1, 60) + ) + report = assess_markdown_fragment(page) + self.assertTrue(report.accepted, [issue.code for issue in report.issues]) + + def test_repetitive_form_rows_are_not_a_loop(self): + # Legal forms legitimately end in identical blank rows. They stay well + # under the 0.35 share, so the page must survive even though its final + # characters do repeat at a fixed period. + preamble = ( + "# Schedule B — Statement of Financial Affairs\n\n" + "The debtor certifies that the following schedule is complete and accurate " + "as of the petition date, and incorporates by reference the exhibits attached " + "hereto in accordance with the applicable local rules of this district. " + ) * 30 + row = " N/A0.00None\n" + form = preamble + "\n\n\n" + (row * 40) + self.assertTrue(assess_markdown_fragment(form).accepted) + # Also with the closing tag, which is what the trailer ladder probes for. + report = assess_markdown_fragment(form + "
") + self.assertTrue(report.accepted, [issue.code for issue in report.issues]) + + def test_sub_threshold_repeating_tail_is_kept(self): + # A tail that repeats but stays under the 0.35 share is not a loop. This + # is also the shape that decides the check's cost: it never short-circuits, + # so it must stay a slice comparison rather than a per-period character walk. + text = (GOOD_SENTENCE * 130) + ("x" * 6_000) + self.assertLess(6_000 / len(text), 0.35) + report = assess_markdown_fragment(text) + self.assertFalse(any(issue.code == "repeated_tail" for issue in report.issues)) + + def test_tail_check_is_cheap_on_a_large_page(self): + # Guards against reintroducing the O(max_period * n) walk, which cost + # ~300ms per page on this input. + text = (GOOD_SENTENCE * 130) + ("x" * 6_000) + start = time.perf_counter() + for _ in range(10): + _has_repeated_tail_loop(text) + self.assertLess((time.perf_counter() - start) / 10, 0.05) + + def test_short_text_does_not_crash(self): + for text in ("", " ", "a", "ab", "# Title"): + assess_markdown_fragment(text) + + +class DisabledChecksTests(unittest.TestCase): + def test_disabling_tail_check_accepts_the_loop(self): + text = (GOOD_SENTENCE * 20) + "\n\n" + (LOOP_SENTENCE * 260) + self.assertFalse(assess_markdown_fragment(text).accepted) + report = assess_markdown_fragment(text, disabled_checks=frozenset({"repeated_tail"})) + self.assertTrue(report.accepted, [issue.code for issue in report.issues]) + + def test_disabling_one_check_leaves_the_others_armed(self): + report = assess_markdown_fragment( + "I'm sorry, but I cannot help with that request.", + disabled_checks=frozenset({"repeated_tail"}), + ) + self.assertFalse(report.accepted) + self.assertTrue(any(issue.code == "refusal_boilerplate" for issue in report.issues)) + + def test_disabling_every_check_accepts_anything(self): + every = frozenset(QUALITY_CHECK_CODES) + for text in ("", "buy now " * 40, "I cannot provide that.", "��� bad"): + self.assertTrue(assess_markdown_fragment(text, disabled_checks=every).accepted) + + def test_verifier_threads_disabled_checks(self): + text = (GOOD_SENTENCE * 20) + "\n\n" + (LOOP_SENTENCE * 260) + self.assertFalse(DeterministicQualityVerifier().classify(text).accepted) + relaxed = DeterministicQualityVerifier(disabled_checks=frozenset({"repeated_tail"})) + self.assertTrue(relaxed.classify(text).accepted) + + def test_all_codes_are_real_checks(self): + # Guards the --disable-quality-check vocabulary against drift in both + # directions: every advertised code must be honoured by a real `enabled(...)` + # guard, and every code the gate raises must be advertised. + self.assertEqual(len(set(QUALITY_CHECK_CODES)), len(QUALITY_CHECK_CODES)) + source = inspect.getsource(assess_markdown_fragment) + for code in QUALITY_CHECK_CODES: + with self.subTest(code=code): + self.assertIn(f'enabled("{code}")', source, f"{code} is advertised but nothing checks it") + for raised in re.findall(r'QualityIssue\("(\w+)"', source): + with self.subTest(raised=raised): + self.assertIn(raised, QUALITY_CHECK_CODES, f"{raised} is raised but not disableable") + + def test_every_code_can_actually_be_disabled(self): + # Each code must be accepted by expand_disabled_checks and land in the set + # the verifier consults; "all" must expand to the whole vocabulary. + for code in QUALITY_CHECK_CODES: + with self.subTest(code=code): + self.assertEqual(expand_disabled_checks([code]), frozenset({code})) + self.assertEqual(expand_disabled_checks(["all"]), frozenset(QUALITY_CHECK_CODES)) + self.assertEqual(expand_disabled_checks([]), frozenset()) + + if __name__ == "__main__": unittest.main()