merge: R-C (.tex ingest) + R-F (section labels) + DQ-5 arXiv fix
This commit is contained in:
Tarik Moussa
@@ -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 is_quality_chunk, section_label
|
||||
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,34 @@ 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: label from the real \section heading — keeps the
|
||||
# controlled bucket when it matches, else stores the cleaned
|
||||
# title (R-F); .txt/.pdf (title=None) classify by content.
|
||||
# NB: `section` semantics now differ by source (write-only col).
|
||||
section_label(title, 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(
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -110,6 +110,46 @@ def classify_section(text: str) -> str:
|
||||
return "body"
|
||||
|
||||
|
||||
def _clean_title(title: str) -> str:
|
||||
"""Normalize a LaTeX ``\\section`` title for use as a section label.
|
||||
|
||||
Strips LaTeX word-commands (``\\emph`` …), accent/escape macros (``\\"o`` →
|
||||
``o``, ``\\&``), ``~`` ties and ``{}$``, drops leading section numbering
|
||||
(``3.2 ``), collapses whitespace, lower-cases, and truncates to 60 chars at a
|
||||
word boundary so the stored ``section`` value is a clean, comparable string.
|
||||
"""
|
||||
t = title.replace("~", " ") # LaTeX non-breaking tie → space
|
||||
t = re.sub(r"\\[a-zA-Z]+\*?", " ", t) # word-commands: \emph, \mathcal …
|
||||
t = re.sub(r"\\[^a-zA-Z]", "", t) # accent/escape macros: \"o → o, \', \`, \&
|
||||
t = re.sub(r"[{}$]", "", t) # braces / math delimiters
|
||||
t = re.sub(r"^\s*\d+(?:\.\d+)*\.?\s*", "", t) # leading numbering "3.2 "
|
||||
t = " ".join(t.split()).strip().lower()
|
||||
return t[:60].rsplit(" ", 1)[0] if len(t) > 60 else t
|
||||
|
||||
|
||||
def section_label(title: str | None, content: str) -> str:
|
||||
"""Label a chunk's section, preferring the real ``\\section`` heading (R-F).
|
||||
|
||||
For section-aware ``.tex`` ingest *title* is the real heading: if it maps to a
|
||||
controlled bucket via :func:`classify_section` (intro / theorem / proof /
|
||||
abstract / bibliography) that bucket is kept (cross-source consistency);
|
||||
otherwise the cleaned title itself is stored (e.g. ``"preliminaries"``,
|
||||
``"rigidity"``) — far more signal than collapsing everything to ``body``.
|
||||
|
||||
When *title* is None / blank (the ``.pdf`` / ``.txt`` / ``run_quality_pass``
|
||||
paths) this falls back to content-based :func:`classify_section` — unchanged
|
||||
behaviour. Safe to store free-text titles because the ``section`` column is
|
||||
write-only (no consumer filters on the controlled vocabulary).
|
||||
"""
|
||||
if not title or not title.strip():
|
||||
return classify_section(content)
|
||||
cleaned = _clean_title(title)
|
||||
if not cleaned:
|
||||
return classify_section(content)
|
||||
bucket = classify_section(cleaned)
|
||||
return bucket if bucket != "body" else cleaned
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Batch filter
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -8,6 +8,7 @@ Provides:
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import gzip
|
||||
import io
|
||||
import logging
|
||||
import tarfile
|
||||
@@ -101,8 +102,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,39 +128,49 @@ 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")
|
||||
except tarfile.TarError as exc:
|
||||
logger.warning("Failed to open tar archive for %s: %s", arxiv_id, exc)
|
||||
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:
|
||||
# Old arXiv submissions are served as a single bare gzipped .tex with no
|
||||
# tar wrapper (e.g. legacy math/* papers) — gunzip the body directly.
|
||||
try:
|
||||
single = gzip.decompress(raw).decode("utf-8", errors="replace")
|
||||
except (OSError, EOFError) as exc:
|
||||
logger.warning("Failed to read arXiv source for %s: %s", arxiv_id, exc)
|
||||
return None
|
||||
if "\\" in single[:5000]: # looks like LaTeX (has backslash commands)
|
||||
return single
|
||||
logger.debug("arXiv source for %s is not LaTeX", arxiv_id)
|
||||
return None
|
||||
|
||||
return None # unreachable but satisfies type checker
|
||||
|
||||
|
||||
def fetch_pdf_url(arxiv_id: str) -> str:
|
||||
"""Return the canonical PDF URL for an arXiv paper.
|
||||
|
||||
@@ -105,6 +105,27 @@ def _normalize_doi(doi: str) -> str:
|
||||
return s
|
||||
|
||||
|
||||
# arXiv works are registered under a DOI of the form 10.48550/arXiv.<id>.
|
||||
_ARXIV_DOI_PREFIX = "10.48550/arxiv."
|
||||
|
||||
|
||||
def _canonical_id(doi: str) -> str:
|
||||
"""Return the canonical ``papers.id`` for an OpenAlex ``doi`` value.
|
||||
|
||||
Builds on :func:`_normalize_doi` (bare, lower-cased DOI), then strips the
|
||||
arXiv DOI prefix ``10.48550/arxiv.`` so an arXiv work's id is the BARE arXiv
|
||||
id (``2305.10988`` / ``math/0603097``) — the form the corpus is keyed on —
|
||||
not the arXiv DOI. Without this, re-ingesting an arXiv paper by its bare id
|
||||
misses ``INSERT … ON CONFLICT (id)`` and trips ``papers_openalex_id_key``.
|
||||
This is the arXiv half of DQ-5 (the DOI half was fixed there; this was
|
||||
explicitly deferred and surfaced again driving the R-C .tex re-ingest).
|
||||
"""
|
||||
normalized = _normalize_doi(doi)
|
||||
if normalized.startswith(_ARXIV_DOI_PREFIX):
|
||||
return normalized[len(_ARXIV_DOI_PREFIX) :]
|
||||
return normalized
|
||||
|
||||
|
||||
def _map_paper(data: dict[str, Any]) -> Paper:
|
||||
authors: list[str] = [
|
||||
a.get("author", {}).get("display_name") or "" for a in data.get("authorships", [])
|
||||
@@ -112,7 +133,7 @@ def _map_paper(data: dict[str, Any]) -> Paper:
|
||||
abstract = _reconstruct_abstract(data.get("abstract_inverted_index"))
|
||||
doi = data.get("doi")
|
||||
return Paper(
|
||||
id=_normalize_doi(doi) if doi else (data.get("id") or ""),
|
||||
id=_canonical_id(doi) if doi else (data.get("id") or ""),
|
||||
title=data.get("title") or "",
|
||||
openalex_id=data.get("id"),
|
||||
authors=authors,
|
||||
|
||||
@@ -396,7 +396,33 @@ is self-contained so a cold session can pick it up.
|
||||
- **Acceptance:** Crossref recovers refs and/or an abstract for ≥1 paper where
|
||||
OpenAlex+S2 are empty.
|
||||
|
||||
### R-C — Prefer arXiv `.tex` source over `.txt` for chunk fidelity · **MED lever, MED-HIGH effort**
|
||||
### R-C — Prefer arXiv `.tex` source over `.txt` for chunk fidelity · **DONE 2026-06-17**
|
||||
|
||||
**Resolution (what was done) — all 13 arXiv papers re-ingested from `.tex`:**
|
||||
- **Multi-file flatten:** `arxiv.fetch_source` now inlines `\input`/`\include`
|
||||
(`tex.flatten_inputs`) so multi-file projects yield the full body, not the
|
||||
primary file's include skeleton (e.g. math/0603097 = 15 files / 12 includes;
|
||||
2305.10988 = 6 / 5). Also handles legacy single-file **bare-gzipped** `.tex`
|
||||
(no tar wrapper) for old math/* papers (commit 121b582).
|
||||
- **Section-aware chunking:** new `tex.chunk_sections` returns `(title, chunk)`
|
||||
pairs (chunks never span a `\section`); the ingest `.tex` path labels the stored
|
||||
`section` from the real heading instead of running `classify_section` on a
|
||||
header-less word-window. Quality-gated as pairs so labels stay aligned.
|
||||
- **Unblocked by the arXiv half of DQ-5** (commit f504ca6): `_canonical_id`
|
||||
strips the `10.48550/arxiv.` DOI prefix so an arXiv paper's id is the bare
|
||||
arXiv id — without it, re-ingest tripped `papers_openalex_id_key` (confirmed
|
||||
live, then fixed; idempotent re-ingest verified).
|
||||
- **Outcome:** `section` column gained signal — ~34/329 arXiv chunks now
|
||||
intro/theorem/proof (was ~all `body`); math/0603097 `{body:31}`→`{body:29,
|
||||
proof:9}`, 2310.17529 gained proof+intro, etc. Math papers with descriptively
|
||||
titled sections stay `body` (classify_section's fixed vocab) — modest but real
|
||||
(R-12). No paper gutted (chunk counts stable/▲); F-16 gate unchanged.
|
||||
- **Files:** `codex/parsing/tex.py`, `codex/sources/arxiv.py`, `codex/ingest.py`,
|
||||
`codex/sources/openalex.py`, `scripts/rc_tex_reingest.py` (driver, dry-run
|
||||
default), + tests. Full suite **370 passed**; ruff/mypy clean. Branch
|
||||
`feat/tex-ingest` (commits 1da81c7, 121b582, f504ca6).
|
||||
|
||||
**Original plan (for context):**
|
||||
- **What:** re-ingest from arXiv LaTeX source through the existing
|
||||
`codex.parsing.tex.latex_to_text` path instead of flattened `.txt`.
|
||||
- **Why:** highest-fidelity input — preserves math, structure, and section
|
||||
@@ -428,6 +454,32 @@ is self-contained so a cold session can pick it up.
|
||||
violates most publisher terms and can get the whole institution's access
|
||||
revoked.** Do not wire a publisher login into the ingest pipeline.
|
||||
|
||||
### R-F — Richer section labels: store real `\section` titles · **DONE 2026-06-17**
|
||||
|
||||
**Resolution (what was done):**
|
||||
- **Why:** R-C labelled `.tex` chunks by mapping the real `\section` heading through
|
||||
`classify_section`'s fixed vocab, so descriptively-titled Math sections collapsed
|
||||
to `body` (only **34/329 = 10 %** non-`body`). A research pass confirmed the
|
||||
`section` column is **write-only** (nothing filters on the vocab — `mcp_server.py`
|
||||
and `wiki.py` don't even SELECT it), so storing free-text titles is safe.
|
||||
- **Change:** new `quality.section_label(title, content)` (+ `_clean_title`) — keep the
|
||||
controlled bucket (intro/theorem/proof/abstract/bibliography) when the heading maps
|
||||
to one, else store the cleaned real title ("the flip algorithm", "main results", …).
|
||||
The `.tex` ingest path uses it; `.pdf`/`.txt`/`run_quality_pass` keep
|
||||
`classify_section` unchanged. `_clean_title` strips LaTeX commands/accents/ties +
|
||||
leading numbering, lower-cases, truncates at a word boundary. ~1 helper + 1 ingest
|
||||
line; no schema change (column already free `TEXT`).
|
||||
- **Result (live):** re-applied to the 13 arXiv papers → non-`body` **34/329 (10 %) →
|
||||
314/329 (95 %)**. Only `math/0306167` (old AMS-TeX, no `\section{}`) stays `body`.
|
||||
- **Tests:** `section_label` / `_clean_title` units + `.tex` ingest stores a
|
||||
descriptive heading verbatim. Full suite 377 passed; ruff/mypy clean. Branch
|
||||
`feat/section-labels` (off `feat/tex-ingest`).
|
||||
- **Accent re-clean — DONE 2026-06-17.** The `_clean_title` polish (`m\"obius`→`mobius`,
|
||||
`~` ties, word-boundary truncation) landed in code, and the existing live labels were
|
||||
re-cleaned in place (apply `_clean_title` to any `chunks.section` containing `\` or
|
||||
`~`; **8 labels** fixed, no re-embed). Verified: **0** labels with LaTeX artifacts
|
||||
remain; non-`body` holds at **314/329 (95 %)**. R-F fully complete.
|
||||
|
||||
---
|
||||
|
||||
## Priority for the next session
|
||||
|
||||
125
scripts/rc_tex_reingest.py
Normal file
125
scripts/rc_tex_reingest.py
Normal 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())
|
||||
@@ -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,20 @@ 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."),
|
||||
("3.2 Discrete Conformal Maps", "Body about discrete conformal maps."),
|
||||
]
|
||||
|
||||
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 +152,12 @@ 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)
|
||||
# R-F: heading → controlled bucket where it maps ("Introduction"→intro,
|
||||
# "Proof of…"→proof), else the cleaned real title ("discrete conformal maps")
|
||||
# instead of collapsing to "body".
|
||||
chunk_rows = chunk_insert_call[0][1]
|
||||
assert [row[4] for row in chunk_rows] == ["intro", "proof", "discrete conformal maps"]
|
||||
assert result.chunks_upserted == len(section_pairs)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -183,7 +193,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 +247,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))
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -6,10 +6,12 @@ from unittest.mock import MagicMock
|
||||
|
||||
from codex.quality import (
|
||||
_bib_score,
|
||||
_clean_title,
|
||||
classify_section,
|
||||
filter_chunks,
|
||||
is_quality_chunk,
|
||||
run_quality_pass,
|
||||
section_label,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -243,3 +245,49 @@ class TestRunQualityPass:
|
||||
kwargs = update_calls[0][0][1]
|
||||
assert kwargs["section"] == "abstract"
|
||||
assert kwargs["id"] == 5
|
||||
|
||||
|
||||
class TestCleanTitle:
|
||||
def test_strips_leading_numbering(self):
|
||||
assert _clean_title("3.2 Variational Principles") == "variational principles"
|
||||
|
||||
def test_strips_latex_and_braces(self):
|
||||
assert _clean_title(r"The \emph{Discrete} Laplacian") == "the discrete laplacian"
|
||||
|
||||
def test_collapses_and_lowercases(self):
|
||||
assert _clean_title(" Rigidity Results ") == "rigidity results"
|
||||
|
||||
def test_strips_accent_macros(self):
|
||||
assert _clean_title(r"M\"obius Invariance") == "mobius invariance"
|
||||
|
||||
def test_strips_ties_and_accents(self):
|
||||
assert _clean_title(r"Colin~de~Verdi\`ere") == "colin de verdiere"
|
||||
|
||||
def test_truncates_at_word_boundary(self):
|
||||
long = "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu omega"
|
||||
out = _clean_title(long)
|
||||
assert len(out) <= 60
|
||||
assert not out.endswith(" ")
|
||||
assert long.lower().startswith(out) # whole words from the start, no partial tail
|
||||
|
||||
|
||||
class TestSectionLabel:
|
||||
def test_title_maps_to_controlled_bucket(self):
|
||||
# Headings matching the controlled vocab keep the canonical bucket.
|
||||
assert section_label("Introduction", "body text") == "intro"
|
||||
assert section_label("Proof of Theorem 3", "body text") == "proof"
|
||||
assert section_label("References", "body text") == "bibliography"
|
||||
|
||||
def test_descriptive_title_stored_verbatim(self):
|
||||
# R-F win: a non-vocab heading is stored as its cleaned title, not "body".
|
||||
assert section_label("Preliminaries", "body text") == "preliminaries"
|
||||
assert section_label("3.1 Discrete Conformal Maps", "x") == "discrete conformal maps"
|
||||
|
||||
def test_none_or_blank_title_falls_back_to_content(self):
|
||||
# .pdf/.txt path: unchanged content-based classification.
|
||||
assert section_label(None, "Theorem 1. Let X be ...") == "theorem"
|
||||
assert section_label("", "ordinary body prose here") == "body"
|
||||
|
||||
def test_title_cleaning_to_empty_falls_back_to_content(self):
|
||||
# A pure-numbering heading cleans to "" → classify by content instead.
|
||||
assert section_label("3.2", "Lemma 2. Suppose ...") == "theorem"
|
||||
|
||||
19
tests/scripts/test_rc_tex_reingest.py
Normal file
19
tests/scripts/test_rc_tex_reingest.py
Normal 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("")
|
||||
@@ -53,6 +53,63 @@ 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
|
||||
|
||||
|
||||
def test_fetch_source_bare_gzip_single_file(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Legacy arXiv submissions are served as a single bare gzipped .tex (no tar
|
||||
wrapper); fetch_source must gunzip the body directly rather than fail."""
|
||||
import gzip
|
||||
|
||||
latex = b"\\documentstyle[a4]{article}\n\\section{Intro}\nOld single-file paper body."
|
||||
gz_bytes = gzip.compress(latex)
|
||||
|
||||
def mock_get(
|
||||
url: str,
|
||||
*,
|
||||
timeout: int = 60,
|
||||
follow_redirects: bool = True,
|
||||
) -> httpx.Response:
|
||||
return httpx.Response(200, content=gz_bytes)
|
||||
|
||||
monkeypatch.setattr(httpx, "get", mock_get)
|
||||
|
||||
result = arxiv.fetch_source("math/0001176")
|
||||
|
||||
assert result is not None
|
||||
assert "Old single-file paper body." in result
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# fetch_source — 404 → None
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -7,7 +7,7 @@ import pytest
|
||||
|
||||
from codex.models import Citation, Paper
|
||||
from codex.sources import openalex
|
||||
from codex.sources.openalex import _normalize_doi, _resolve_id
|
||||
from codex.sources.openalex import _canonical_id, _normalize_doi, _resolve_id
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
@@ -97,6 +97,42 @@ def test_normalize_doi_already_bare_is_idempotent() -> None:
|
||||
assert _normalize_doi("10.1007/s00454-019-00132-8") == "10.1007/s00454-019-00132-8"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _canonical_id — arXiv DOI → bare arXiv id (arXiv half of DQ-5)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_canonical_id_strips_arxiv_doi_modern() -> None:
|
||||
# arXiv works carry the arXiv DOI; the canonical id is the bare arXiv id.
|
||||
assert _canonical_id("https://doi.org/10.48550/arXiv.2305.10988") == "2305.10988"
|
||||
|
||||
|
||||
def test_canonical_id_strips_arxiv_doi_legacy() -> None:
|
||||
assert _canonical_id("https://doi.org/10.48550/arXiv.math/0603097") == "math/0603097"
|
||||
|
||||
|
||||
def test_canonical_id_leaves_regular_doi_bare() -> None:
|
||||
# A normal journal DOI is only normalized (bare + lower-case), not stripped.
|
||||
assert _canonical_id("https://doi.org/10.1007/S00454-019-00132-8") == (
|
||||
"10.1007/s00454-019-00132-8"
|
||||
)
|
||||
|
||||
|
||||
def test_fetch_paper_arxiv_doi_yields_bare_id(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""An arXiv work maps to the BARE arXiv id (not the arXiv DOI), so re-ingesting
|
||||
by arXiv id hits ON CONFLICT (id) instead of tripping papers_openalex_id_key."""
|
||||
work = {**_SAMPLE_WORK, "doi": "https://doi.org/10.48550/arXiv.2305.10988"}
|
||||
|
||||
def mock_get(url: str, params: dict[str, str] | None = None) -> httpx.Response:
|
||||
return httpx.Response(200, json=work)
|
||||
|
||||
monkeypatch.setattr(openalex, "_get", mock_get)
|
||||
paper = openalex.fetch_paper("2305.10988")
|
||||
|
||||
assert paper is not None
|
||||
assert paper.id == "2305.10988"
|
||||
|
||||
|
||||
def test_map_paper_falls_back_to_openalex_id_without_doi(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
|
||||
Reference in New Issue
Block a user