diff --git a/tests/test_latex/conftest.py b/tests/test_latex/conftest.py index 60c5429b71..64e74fecfd 100644 --- a/tests/test_latex/conftest.py +++ b/tests/test_latex/conftest.py @@ -37,6 +37,26 @@ def pytest_addoption(parser: pytest.Parser) -> None: action="store_true", help="Save successful arXiv LaTeX sources and DocLang output.", ) + parser.addoption( + "--recheck-failures", + action="store_true", + help="Convert again the arXiv LaTeX files that failed in an earlier run.", + ) + parser.addoption( + "--corpus-shard", + action="append", + type=int, + default=None, + help="Read only this shard (0-79) of the arXiv LaTeX corpus. Repeat it to " + "read more shards.", + ) + parser.addoption( + "--corpus-dir", + type=Path, + default=None, + help="Write the arXiv LaTeX corpus manifest, report, and saved files to " + "this folder instead of the repository root.", + ) @pytest.fixture(scope="module") diff --git a/tests/test_latex/test_large_scale.py b/tests/test_latex/test_large_scale.py index 4d9fc32a15..6dd8ba0919 100644 --- a/tests/test_latex/test_large_scale.py +++ b/tests/test_latex/test_large_scale.py @@ -10,12 +10,21 @@ The full stream contains millions of source files and can take a long time. Completed files are appended to ``latex-manifest.jsonl``. Run the same command again after an interruption to continue with the next unprocessed source file. +Add ``--recheck-failures`` to convert the files that failed earlier again, for +example after a fix. The latest result of each file replaces the earlier one. +``latex-failures.jsonl`` keeps every failure found, also failures that a later +recheck cleared. +The stream starts with shard 0, which holds the arXiv papers of 2007 and 2008. +Use ``--corpus-shard N`` (0-79, repeatable) to read other shards, and +``--corpus-dir DIR`` to write the output files outside the repository, for +example one folder for each of several parallel runs. """ import hashlib import json import logging from collections import Counter +from collections.abc import Iterable from io import BytesIO from pathlib import Path from xml.etree import ElementTree as ET @@ -36,9 +45,16 @@ ) _DATASET = "secemp9/arxiv-complete" +_SHARD_COUNT = 80 _SOURCE_EXTENSIONS = {"tex", "ltx"} +def _shard_files(shards: list[int]) -> list[str]: + return [ + f"latex/train-{shard:05d}-of-{_SHARD_COUNT:05d}.parquet" for shard in shards + ] + + class _ManifestEntry(BaseModel): paper_id: str path: str @@ -72,6 +88,14 @@ def _read_manifest(path: Path) -> dict[tuple[str, str], _ManifestEntry]: return entries +def _needs_check(entry: _ManifestEntry | None, recheck_failures: bool) -> bool: + return entry is None or (recheck_failures and bool(entry.stages)) + + +def _failure_counts(entries: Iterable[_ManifestEntry]) -> Counter[str]: + return Counter(stage for entry in entries for stage in entry.stages) + + def test_manifest_recovers_interrupted_write(tmp_path: Path) -> None: manifest_path = tmp_path / "latex-manifest.jsonl" entry = _ManifestEntry(paper_id="1234.5678", path="main.tex", stages=[]) @@ -82,6 +106,34 @@ def test_manifest_recovers_interrupted_write(tmp_path: Path) -> None: assert manifest_path.read_bytes() == complete_line +def test_recheck_replaces_earlier_failure(tmp_path: Path) -> None: + """A file that passes on a recheck no longer counts as a corpus failure.""" + failed = _ManifestEntry(paper_id="1234.5678", path="main.tex", stages=["parse"]) + passed = _ManifestEntry(paper_id="1234.5678", path="main.tex", stages=[]) + other = _ManifestEntry(paper_id="1234.5679", path="main.tex", stages=[]) + + assert _needs_check(None, recheck_failures=False) + assert not _needs_check(failed, recheck_failures=False) + assert _needs_check(failed, recheck_failures=True) + assert not _needs_check(other, recheck_failures=True) + + manifest_path = tmp_path / "latex-manifest.jsonl" + manifest_path.write_text( + "".join(entry.model_dump_json() + "\n" for entry in (failed, other, passed)), + encoding="utf-8", + ) + completed = _read_manifest(manifest_path) + assert completed[("1234.5678", "main.tex")] == passed + assert not _failure_counts(completed.values()) + + +def test_shard_files_name_the_latex_parquet_files() -> None: + assert _shard_files([0, 79]) == [ + "latex/train-00000-of-00080.parquet", + "latex/train-00079-of-00080.parquet", + ] + + def _local_name(tag: str) -> str: return tag.rsplit("}", 1)[-1] @@ -119,28 +171,31 @@ def test_arxiv_latex_corpus( if max_items is not None and max_items <= 0: pytest.fail("--max-items must be a positive integer") write_success: bool = pytestconfig.getoption("--write-success") + recheck_failures: bool = pytestconfig.getoption("--recheck-failures") + shards: list[int] = pytestconfig.getoption("--corpus-shard") or [] + if any(not 0 <= shard < _SHARD_COUNT for shard in shards): + pytest.fail(f"--corpus-shard must be between 0 and {_SHARD_COUNT - 1}") + corpus_dir: Path | None = pytestconfig.getoption("--corpus-dir") from datasets import load_dataset from doclang import ValidationError, validate - repo_root = Path(__file__).resolve().parents[2] - manifest_path = repo_root / "latex-manifest.jsonl" + output_dir = corpus_dir or Path(__file__).resolve().parents[2] + output_dir.mkdir(parents=True, exist_ok=True) + manifest_path = output_dir / "latex-manifest.jsonl" completed = _read_manifest(manifest_path) - report_path = repo_root / "latex-failures.jsonl" + report_path = output_dir / "latex-failures.jsonl" if completed: report_path.touch(exist_ok=True) else: report_path.write_text("", encoding="utf-8") dclg_path = tmp_path / "current.dclg" - artifacts_dir = repo_root / "latex-failures" + artifacts_dir = output_dir / "latex-failures" source_dir = artifacts_dir / "sources" invalid_dir = artifacts_dir / "invalid-dclg" - success_dir = repo_root / "latex-success" + success_dir = output_dir / "latex-success" if write_success: success_dir.mkdir(exist_ok=True) - failures: Counter[str] = Counter( - stage for entry in completed.values() for stage in entry.stages - ) checked = 0 row_stages: list[str] = [] @@ -171,17 +226,23 @@ def report(row: dict, stage: str, error: str, dclg: str | None = None) -> None: entry["dclg"] = str(artifact) with report_path.open("a", encoding="utf-8") as output: output.write(json.dumps(entry, ensure_ascii=False) + "\n") - failures[stage] += 1 row_stages.append(stage) - corpus = load_dataset(_DATASET, "latex", split="train", streaming=True) + if shards: + corpus = load_dataset( + _DATASET, data_files=_shard_files(shards), split="train", streaming=True + ) + else: + corpus = load_dataset(_DATASET, "latex", split="train", streaming=True) with caplog.at_level(logging.WARNING, logger="docling.backend.latex"): for row in corpus: if max_items is not None and checked >= max_items: break if row["ext"] not in _SOURCE_EXTENSIONS: continue - if (row["paper_id"], row["path"]) in completed: + if not _needs_check( + completed.get((row["paper_id"], row["path"])), recheck_failures + ): continue checked += 1 row_stages = [] @@ -296,7 +357,7 @@ def report(row: dict, stage: str, error: str, dclg: str | None = None) -> None: finish_row(row) summary = ( - f"Checked {checked} new source files ({len(completed)} total); " + f"Checked {checked} source files ({len(completed)} in the manifest); " f"manifest: {manifest_path}; report: {report_path}" ) terminal_reporter = pytestconfig.pluginmanager.get_plugin("terminalreporter") @@ -305,4 +366,5 @@ def report(row: dict, stage: str, error: str, dclg: str | None = None) -> None: else: terminal_reporter.write_line(summary) assert completed, "The LaTeX stream contained no .tex or .ltx files" + failures = _failure_counts(completed.values()) assert not failures, f"Corpus failures: {dict(failures)}; report: {report_path}"