From 1da81c7b28edf9ea05be544735aa3b63ed364574 Mon Sep 17 00:00:00 2001 From: Tarik Moussa Date: Thu, 18 Jun 2026 04:15:20 +0200 Subject: [PATCH] feat(tex): section-aware multi-file .tex ingest (R-C) Flatten multi-file arXiv LaTeX (\input/\include) in arxiv.fetch_source so the full body is assembled rather than just the primary file's include skeleton, and add section-aware chunking (tex.chunk_sections) so .tex chunks never span a \section boundary and are labelled by their real heading. The ingest .tex path now produces (section, chunk) pairs, quality-gated together so labels stay aligned; the stored 'section' column gains genuine signal instead of mostly 'body' (serves DQ-3 fidelity + audit R-12). Adds scripts/rc_tex_reingest.py to re-ingest arXiv papers from .tex (dry-run by default; replaces chunks). Tests: flatten_inputs, chunk_sections, multi-file fetch_source, section-label ingest, is_arxiv_id. Full suite 365 passed; ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 --- codex/ingest.py | 35 +++++--- codex/parsing/tex.py | 60 +++++++++++++ codex/sources/arxiv.py | 48 +++++----- scripts/rc_tex_reingest.py | 125 ++++++++++++++++++++++++++ tests/ingest/test_ingest.py | 23 +++-- tests/parsing/test_tex.py | 55 +++++++++++- tests/scripts/test_rc_tex_reingest.py | 19 ++++ tests/sources/test_arxiv.py | 33 +++++++ 8 files changed, 356 insertions(+), 42 deletions(-) create mode 100644 scripts/rc_tex_reingest.py create mode 100644 tests/scripts/test_rc_tex_reingest.py diff --git a/codex/ingest.py b/codex/ingest.py index c0f768f..ab79250 100644 --- a/codex/ingest.py +++ b/codex/ingest.py @@ -12,7 +12,7 @@ from codex.config import get_settings from codex.db import get_conn from codex.embed import get_embedder from codex.models import Citation, Paper -from codex.quality import classify_section, filter_chunks +from codex.quality import classify_section, is_quality_chunk from codex.sources import arxiv, crossref, openalex, semanticscholar logger = logging.getLogger(__name__) @@ -262,17 +262,21 @@ def ingest_paper( if source_path is not None: from codex.parsing.grobid import extract_references from codex.parsing.nougat import pdf_to_markdown - from codex.parsing.tex import chunk_text, latex_to_text + from codex.parsing.tex import chunk_sections, chunk_text suffix = Path(source_path).suffix.lower() - text = "" + # (section_title | None, chunk_text) pairs. `.tex` is section-aware so + # the stored `section` column reflects the real heading (R-C / R-12); + # `.txt`/`.pdf` carry no section structure, so the title is None. + raw_pairs: list[tuple[str | None, str]] = [] if suffix == ".tex": - text = latex_to_text(Path(source_path).read_text()) + raw_pairs = chunk_sections(Path(source_path).read_text()) elif suffix == ".txt": text = Path(source_path).read_text(encoding="utf-8", errors="replace") + raw_pairs = [(None, c) for c in chunk_text(text)] elif suffix == ".pdf": - text = pdf_to_markdown(source_path) + raw_pairs = [(None, c) for c in chunk_text(pdf_to_markdown(source_path))] # Also extract GROBID refs grobid_refs = extract_references(source_path) for ref in grobid_refs: @@ -287,24 +291,33 @@ def ingest_paper( else: logger.warning("Unbekannter Dateityp: %s — kein Text-Parsing", source_path) - raw_chunks = chunk_text(text) if text else [] - chunks_text = filter_chunks(raw_chunks, settings=get_settings()) + # Quality-gate the pairs (keeps each chunk aligned with its section title). + settings = get_settings() + kept_pairs = [ + (title, content) + for title, content in raw_pairs + if is_quality_chunk(content, settings=settings) + ] # Delete existing chunks for this paper conn.execute("DELETE FROM chunks WHERE paper_id = %s", (paper.id,)) - if chunks_text: + if kept_pairs: # Embed all chunks in one batch - chunk_embeddings = embedder.encode_dense(chunks_text) + chunk_contents = [content for _, content in kept_pairs] + chunk_embeddings = embedder.encode_dense(chunk_contents) chunk_rows = [ ( paper.id, ord_idx, content, chunk_embeddings[ord_idx].tolist(), - classify_section(content), + # .tex: classify from the section heading (prepended to the + # snippet) so the label reflects the real section; otherwise + # fall back to content-only classification. + classify_section(f"{title}. {content}" if title else content), ) - for ord_idx, content in enumerate(chunks_text) + for ord_idx, (title, content) in enumerate(kept_pairs) ] with conn.cursor() as cur: cur.executemany( diff --git a/codex/parsing/tex.py b/codex/parsing/tex.py index 43514d6..ed77ad7 100644 --- a/codex/parsing/tex.py +++ b/codex/parsing/tex.py @@ -17,6 +17,10 @@ _SECTION_RE = re.compile( re.DOTALL, ) +# \input{file} / \include{file} — arXiv papers routinely split the body across +# files, so these must be inlined before parsing (R-C multi-file flattening). +_INPUT_RE = re.compile(r"\\(?:input|include)\s*\{([^}]+)\}") + # Patterns for LaTeX noise removal (applied in order). _COMMENT_RE = re.compile(r"%[^\n]*") _CITE_RE = re.compile(r"\\cite\*?\{[^}]*\}") @@ -157,3 +161,59 @@ def latex_to_text(latex: str) -> str: return "\n\n".join(body for _, body in sections) # Fallback: no section structure — clean the whole string. return _clean_latex(latex) + + +def _norm_texkey(name: str) -> str: + """Normalize a tex file name / ``\\input`` argument to a comparable key. + + Strips a leading ``./``, any directory-irrelevant whitespace, and a trailing + ``.tex`` so ``\\input{sections/intro}`` matches the archive member + ``sections/intro.tex``. + """ + n = name.strip().lstrip("./") + return n[:-4] if n.endswith(".tex") else n + + +def flatten_inputs(primary: str, files: dict[str, str]) -> str: + """Recursively inline ``\\input``/``\\include`` directives (R-C). + + *files* maps each source file's archive name to its contents (keys may carry + a ``.tex`` suffix or a ``./`` prefix — they are normalized internally). arXiv + multi-file projects leave the ``\\documentclass`` file as little more than a + skeleton of ``\\input`` lines, so the body must be assembled before parsing. + Unresolved references are dropped; recursion is capped at depth 10 to guard + against include cycles. + """ + norm = {_norm_texkey(k): v for k, v in files.items()} + + def _expand(text: str, depth: int) -> str: + if depth > 10: + return text + + def _repl(match: re.Match[str]) -> str: + content = norm.get(_norm_texkey(match.group(1))) + return _expand(content, depth + 1) if content is not None else "" + + return _INPUT_RE.sub(_repl, text) + + return _expand(primary, 0) + + +def chunk_sections(latex: str, size: int = 512, overlap: int = 64) -> list[tuple[str | None, str]]: + """Chunk LaTeX *within* each section, pairing every chunk with its title. + + Unlike :func:`latex_to_text` (which discards titles and flattens the whole + document into one block before word-window chunking), this keeps chunks from + spanning section boundaries and remembers which ``\\section`` each came from. + The caller can then label the stored ``section`` column from the real heading + instead of mostly ``body`` (R-C / audit R-12). Falls back to title-less + chunks of the whole cleaned document when there are no section commands. + """ + sections = extract_sections(latex) + if not sections: + return [(None, chunk) for chunk in chunk_text(_clean_latex(latex), size, overlap)] + pairs: list[tuple[str | None, str]] = [] + for title, body in sections: + for chunk in chunk_text(body, size, overlap): + pairs.append((title or None, chunk)) + return pairs diff --git a/codex/sources/arxiv.py b/codex/sources/arxiv.py index 9bb0303..dd61961 100644 --- a/codex/sources/arxiv.py +++ b/codex/sources/arxiv.py @@ -101,8 +101,9 @@ def fetch_source(arxiv_id: str) -> str | None: Downloads the .tar.gz source bundle from ``https://arxiv.org/src/{arxiv_id}``, locates the primary .tex file (preferring any file containing ``\\documentclass``, - falling back to the largest .tex by size), and returns its contents as a UTF-8 - string. + falling back to the largest .tex by size), and inlines its ``\\input``/ + ``\\include`` directives from the other archive members so multi-file projects + return the full document body, not just the primary file's include skeleton. Parameters ---------- @@ -126,33 +127,36 @@ def fetch_source(arxiv_id: str) -> str | None: if response.status_code != 200: response.raise_for_status() + from codex.parsing.tex import flatten_inputs + raw = response.content try: with tarfile.open(fileobj=io.BytesIO(raw), mode="r:gz") as tf: - tex_members = [m for m in tf.getmembers() if m.name.endswith(".tex")] - if not tex_members: - logger.debug("No .tex files found in arXiv source for %s", arxiv_id) - return None - - # Prefer the file containing \documentclass (primary document) - primary: tarfile.TarInfo | None = None - for member in tex_members: + files: dict[str, str] = {} + primary_name: str | None = None + for member in tf.getmembers(): + if not member.name.endswith(".tex"): + continue f = tf.extractfile(member) if f is None: continue - content_bytes = f.read() - if b"\\documentclass" in content_bytes: - primary = member - # Decode and return immediately — first match wins - return content_bytes.decode("utf-8", errors="replace") + content = f.read().decode("utf-8", errors="replace") + files[member.name] = content + # Prefer the \documentclass file as the primary document. + if primary_name is None and "\\documentclass" in content: + primary_name = member.name - if primary is None: - # Fallback: largest .tex by size - largest = max(tex_members, key=lambda m: m.size) - f = tf.extractfile(largest) - if f is None: - return None - return f.read().decode("utf-8", errors="replace") + if not files: + logger.debug("No .tex files found in arXiv source for %s", arxiv_id) + return None + + # No \documentclass anywhere → fall back to the largest .tex. + if primary_name is None: + primary_name = max(files, key=lambda k: len(files[k])) + + # Inline \input/\include so multi-file projects yield the full body, + # not just the primary file's skeleton of include directives (R-C). + return flatten_inputs(files[primary_name], files) except tarfile.TarError as exc: logger.warning("Failed to open tar archive for %s: %s", arxiv_id, exc) return None diff --git a/scripts/rc_tex_reingest.py b/scripts/rc_tex_reingest.py new file mode 100644 index 0000000..53839ac --- /dev/null +++ b/scripts/rc_tex_reingest.py @@ -0,0 +1,125 @@ +"""R-C — re-ingest arXiv papers from LaTeX source (section-aware, higher fidelity). + +For each arXiv paper, download its source via :func:`codex.sources.arxiv.fetch_source` +(multi-file ``\\input``/``\\include`` flattened), write it to a temp ``.tex``, and +re-ingest through :func:`codex.ingest.ingest_paper`. The ``.tex`` ingest path is +section-aware (:func:`codex.parsing.tex.chunk_sections`), so the stored ``section`` +column is labelled from real ``\\section`` headings instead of mostly ``body`` — +serving DQ-3 (fidelity vs the OCR'd ``.txt``) and audit R-12 (weak ``section`` column). + +Only arXiv papers have downloadable source; DOI/journal papers and books are skipped. +Re-ingest REPLACES the paper's chunks (DELETE + re-INSERT) — that is the intended R-C +improvement (flattened ``.txt`` → clean ``.tex`` chunks). It is idempotent and +**dry-run by default**; pass ``--write`` to mutate the live corpus. + +Usage +----- + # dry-run: list arXiv papers + source availability (no writes) + PYTHONPATH=. python scripts/rc_tex_reingest.py + + # re-ingest every arXiv paper from .tex + PYTHONPATH=. python scripts/rc_tex_reingest.py --write + + # one paper + PYTHONPATH=. python scripts/rc_tex_reingest.py --only 2305.10988 --write + +Needs ``DATABASE_URL`` (Jetson tunnel) as documented in +``docs/audit/DATA-QUALITY-2026-06-15.md``. +""" + +from __future__ import annotations + +import argparse +import logging +import tempfile + +import psycopg +import psycopg.rows + +from codex.db import get_conn +from codex.ingest import ingest_paper +from codex.sources import arxiv + +logger = logging.getLogger(__name__) + + +def is_arxiv_id(paper_id: str) -> bool: + """Return True for arXiv ids (the only papers with downloadable LaTeX source). + + arXiv ids are modern (``2305.10988``) or legacy (``math/0603097``). DOIs start + with ``10.``; OpenAlex W-ids start with ``W`` or an ``https://openalex.org/`` URL. + """ + p = (paper_id or "").strip() + return bool(p) and not p.startswith("10.") and not p.startswith(("W", "https://")) + + +def _section_dist(conn: psycopg.Connection[psycopg.rows.DictRow], paper_id: str) -> dict[str, int]: + """Return the section-label distribution for a paper's stored chunks.""" + rows = conn.execute( + "SELECT coalesce(section, '(null)') AS s, count(*) AS n " + "FROM chunks WHERE paper_id = %s GROUP BY 1 ORDER BY 2 DESC", + (paper_id,), + ).fetchall() + return {r["s"]: int(r["n"]) for r in rows} + + +def main(argv: list[str] | None = None) -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--only", default=None, help="restrict to a single arXiv id") + parser.add_argument("--limit", type=int, default=None, help="cap number of papers") + parser.add_argument( + "--write", action="store_true", help="re-ingest into the live DB (default: dry-run)" + ) + args = parser.parse_args(argv) + logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s") + + with get_conn() as conn: + rows = conn.execute("SELECT id FROM papers ORDER BY id").fetchall() + targets = [ + r["id"] + for r in rows + if is_arxiv_id(r["id"]) and (args.only is None or r["id"] == args.only) + ] + if args.limit is not None: + targets = targets[: args.limit] + + logger.info( + "arXiv papers to process=%d | mode=%s", + len(targets), + "WRITE" if args.write else "DRY-RUN", + ) + + for pid in targets: + try: + src = arxiv.fetch_source(pid) + except Exception as exc: + logger.warning("%s: source fetch failed: %s", pid, exc) + continue + if not src: + logger.warning("%s: no LaTeX source available — skipping", pid) + continue + if not args.write: + logger.info("%s: source OK (%d chars)", pid, len(src)) + continue + + with get_conn() as conn: + before = _section_dist(conn, pid) + with tempfile.NamedTemporaryFile("w", suffix=".tex", delete=True) as tf: + tf.write(src) + tf.flush() + res = ingest_paper(pid, source_path=tf.name) + with get_conn() as conn: + after = _section_dist(conn, pid) + logger.info( + "%s: chunks %d->%d | section before=%s after=%s", + pid, + sum(before.values()), + res.chunks_upserted, + before, + after, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/ingest/test_ingest.py b/tests/ingest/test_ingest.py index 936e6cb..8828419 100644 --- a/tests/ingest/test_ingest.py +++ b/tests/ingest/test_ingest.py @@ -113,7 +113,8 @@ def test_ingest_paper_basic() -> None: def test_ingest_paper_source_tex(tmp_path: Any) -> None: - """``.tex`` source → latex_to_text, chunk_text, embed, chunks inserted in DB.""" + """``.tex`` source → section-aware chunking; the stored ``section`` column is + labelled from the real ``\\section`` heading, not the chunk body (R-C / R-12).""" tex_file = tmp_path / "paper.tex" tex_file.write_text(r"\section{Introduction} Hello world. This is a test paper.") @@ -123,16 +124,19 @@ def test_ingest_paper_source_tex(tmp_path: Any) -> None: mock_conn.executemany = MagicMock() mock_conn.commit = MagicMock() - fake_chunks = ["chunk one text here.", "chunk two text here."] + # (section title, chunk content) pairs from the section-aware chunker. + section_pairs = [ + ("Introduction", "Body of the introduction goes here."), + ("Proof of the Main Theorem", "Body of the proof goes here."), + ] with ( patch("codex.ingest.openalex.fetch_paper", return_value=paper), patch("codex.ingest.openalex.fetch_citations", return_value=[]), patch("codex.ingest.get_embedder", return_value=_fake_embedder()), patch("codex.ingest.get_conn", side_effect=_make_conn_cm(mock_conn)), - patch("codex.parsing.tex.latex_to_text", return_value="Hello world. Intro text."), - patch("codex.parsing.tex.chunk_text", return_value=fake_chunks), - patch("codex.ingest.filter_chunks", side_effect=lambda chunks, settings: chunks), + patch("codex.parsing.tex.chunk_sections", return_value=section_pairs), + patch("codex.ingest.is_quality_chunk", return_value=True), ): result = ingest_paper(paper.id, source_path=str(tex_file)) @@ -147,7 +151,10 @@ def test_ingest_paper_source_tex(tmp_path: Any) -> None: chunk_sql: str = chunk_insert_call[0][0] assert "INSERT INTO chunks" in chunk_sql - assert result.chunks_upserted == len(fake_chunks) + # Section label comes from the heading: "Introduction" → intro, "Proof of…" → proof. + chunk_rows = chunk_insert_call[0][1] + assert [row[4] for row in chunk_rows] == ["intro", "proof"] + assert result.chunks_upserted == len(section_pairs) # --------------------------------------------------------------------------- @@ -183,7 +190,7 @@ def test_ingest_paper_source_pdf(tmp_path: Any) -> None: patch("codex.parsing.nougat.pdf_to_markdown", return_value="pdf markdown text."), patch("codex.parsing.tex.chunk_text", return_value=fake_chunks), patch("codex.parsing.grobid.extract_references", return_value=grobid_refs), - patch("codex.ingest.filter_chunks", side_effect=lambda chunks, settings: chunks), + patch("codex.ingest.is_quality_chunk", return_value=True), ): result = ingest_paper(paper.id, source_path=str(pdf_file)) @@ -237,7 +244,7 @@ def test_ingest_pdf_grobid_citations_normalize_doi(tmp_path: Any) -> None: patch("codex.parsing.nougat.pdf_to_markdown", return_value="pdf markdown text."), patch("codex.parsing.tex.chunk_text", return_value=["pdf chunk one."]), patch("codex.parsing.grobid.extract_references", return_value=grobid_refs), - patch("codex.ingest.filter_chunks", side_effect=lambda chunks, settings: chunks), + patch("codex.ingest.is_quality_chunk", return_value=True), ): result = ingest_paper(paper.id, source_path=str(pdf_file)) diff --git a/tests/parsing/test_tex.py b/tests/parsing/test_tex.py index 002cc23..13a8044 100644 --- a/tests/parsing/test_tex.py +++ b/tests/parsing/test_tex.py @@ -4,7 +4,13 @@ from __future__ import annotations import pytest -from codex.parsing.tex import chunk_text, extract_sections, latex_to_text +from codex.parsing.tex import ( + chunk_sections, + chunk_text, + extract_sections, + flatten_inputs, + latex_to_text, +) class TestExtractSections: @@ -110,3 +116,50 @@ class TestLatexToText: result = latex_to_text(latex) assert r"\cite" not in result assert "%" not in result + + +class TestFlattenInputs: + def test_inlines_input_and_include(self) -> None: + files = { + "main.tex": r"Start \input{sec/intro} mid \include{proof} end", + "sec/intro.tex": "INTRO BODY", + "proof.tex": "PROOF BODY", + } + out = flatten_inputs(files["main.tex"], files) + assert "INTRO BODY" in out + assert "PROOF BODY" in out + assert r"\input" not in out and r"\include" not in out + + def test_resolves_without_tex_suffix(self) -> None: + # \input{intro} must resolve the archive member "intro.tex". + files = {"main.tex": r"\input{intro}", "intro.tex": "BODY"} + assert "BODY" in flatten_inputs(files["main.tex"], files) + + def test_recursive_includes(self) -> None: + files = {"a.tex": r"X \input{b}", "b.tex": r"Y \input{c}", "c.tex": "Z"} + out = flatten_inputs(files["a.tex"], files) + assert "X" in out and "Y" in out and "Z" in out + + def test_unresolved_input_dropped(self) -> None: + # A reference with no matching file is removed, not left as a directive. + out = flatten_inputs(r"A \input{missing} B", {}) + assert r"\input" not in out + assert "A" in out and "B" in out + + +class TestChunkSections: + def test_pairs_carry_section_titles(self) -> None: + latex = ( + r"\section{Introduction}" + "\nIntro body text here.\n" + r"\section{Proof}" + "\nProof body text here.\n" + ) + pairs = chunk_sections(latex) + titles = [t for t, _ in pairs] + assert "Introduction" in titles + assert "Proof" in titles + # No chunk spans a section boundary: each chunk belongs to one title. + assert all(isinstance(c, str) and c for _, c in pairs) + + def test_fallback_titleless_when_no_sections(self) -> None: + pairs = chunk_sections("Plain text with no section commands at all here.") + assert pairs and all(title is None for title, _ in pairs) diff --git a/tests/scripts/test_rc_tex_reingest.py b/tests/scripts/test_rc_tex_reingest.py new file mode 100644 index 0000000..17a8770 --- /dev/null +++ b/tests/scripts/test_rc_tex_reingest.py @@ -0,0 +1,19 @@ +"""Tests for scripts.rc_tex_reingest.""" + +from __future__ import annotations + +from scripts.rc_tex_reingest import is_arxiv_id + + +def test_is_arxiv_id_accepts_modern_and_legacy() -> None: + assert is_arxiv_id("2305.10988") + assert is_arxiv_id("math/0603097") + assert is_arxiv_id("1005.2698") + + +def test_is_arxiv_id_rejects_dois_and_wids() -> None: + assert not is_arxiv_id("10.1007/s00454-019-00132-8") + assert not is_arxiv_id("10.14279/depositonce-20357") + assert not is_arxiv_id("W2971636899") + assert not is_arxiv_id("https://openalex.org/W2971636899") + assert not is_arxiv_id("") diff --git a/tests/sources/test_arxiv.py b/tests/sources/test_arxiv.py index c4ceb60..f62ff82 100644 --- a/tests/sources/test_arxiv.py +++ b/tests/sources/test_arxiv.py @@ -53,6 +53,39 @@ def test_fetch_source_returns_latex(monkeypatch: pytest.MonkeyPatch) -> None: assert "Hello world" in result +def test_fetch_source_flattens_multifile(monkeypatch: pytest.MonkeyPatch) -> None: + """Multi-file projects: \\input/\\include directives are inlined so the body + from the included files is returned, not just the primary skeleton (R-C).""" + main = ( + b"\\documentclass{article}\\begin{document}" + b"\\input{sections/intro}\\include{proof}\\end{document}" + ) + tar_bytes = _make_tar_gz( + { + "main.tex": main, + "sections/intro.tex": b"INTRODUCTION BODY TEXT", + "proof.tex": b"PROOF BODY TEXT", + } + ) + + def mock_get( + url: str, + *, + timeout: int = 60, + follow_redirects: bool = True, + ) -> httpx.Response: + return httpx.Response(200, content=tar_bytes) + + monkeypatch.setattr(httpx, "get", mock_get) + + result = arxiv.fetch_source("2301.07041") + + assert result is not None + assert "INTRODUCTION BODY TEXT" in result # \input was inlined + assert "PROOF BODY TEXT" in result # \include was inlined + assert "\\input" not in result and "\\include" not in result + + # --------------------------------------------------------------------------- # fetch_source — 404 → None # ---------------------------------------------------------------------------