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 <noreply@anthropic.com>
This commit is contained in:
Tarik Moussa
2026-06-18 04:15:20 +02:00
parent 1516684bbb
commit 1da81c7b28
8 changed files with 356 additions and 42 deletions

View File

@@ -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(

View File

@@ -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

View File

@@ -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

125
scripts/rc_tex_reingest.py Normal file
View File

@@ -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())

View File

@@ -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))

View File

@@ -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)

View File

@@ -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("")

View File

@@ -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
# ---------------------------------------------------------------------------