feat(F-09): rich parsing — formula + figure extraction
- codex/parsing/mathpix.py: pix2tex (local, CPU) primary + MathPix API optional; bbox heuristic h>15px, math-char-count>5; singleton model cache - codex/parsing/figures.py: pymupdf embedded-image extraction → PNG; caption detection via proximity + "Figure/Fig./Abbildung" prefix - codex/models.py: FormulaChunk + FigureChunk dataclasses (R-10/R-11) - codex/ingest.py: --rich flag wires formula+figure extraction post-ingest - codex/cli.py: search_app sub-typer (paper + formula subcommands), --rich flag on ingest; wiki_app from F-12 preserved intact - codex/config.py: mathpix_app_id/key, pix2tex_fallback, figures_dir - infra/schema.sql: formulas + figures tables with HNSW pgvector indexes - pyproject.toml: pymupdf>=1.24, pix2tex>=0.1.4 - tests/parsing/test_mathpix.py + test_figures.py: 31 tests (mock pix2tex + MathPix HTTP, real pymupdf on synthetic PDF) Gate: 158 passed, ruff clean, mypy clean (20 files) Requirements: R-10 R-11 R-12 R-13 R-14 → done Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
Tarik Moussa
311
codex/parsing/mathpix.py
Normal file
311
codex/parsing/mathpix.py
Normal file
@@ -0,0 +1,311 @@
|
||||
"""Formula extraction from PDF pages.
|
||||
|
||||
Two backends are supported:
|
||||
|
||||
1. **MathPix** (cloud): POST the PDF to ``api.mathpix.com/v3/pdf`` and poll
|
||||
for the result JSON. Requires ``MATHPIX_APP_ID`` and ``MATHPIX_APP_KEY``
|
||||
environment variables (or ``.env`` file).
|
||||
|
||||
2. **pix2tex** (local, MIT licence): crop each detected math bounding-box at
|
||||
3× resolution and run the ``LatexOCR`` model locally. Used when MathPix
|
||||
credentials are absent. Model (~116 MB) is downloaded on first call.
|
||||
|
||||
The public entry point :func:`extract_formulas` routes to whichever backend
|
||||
is available.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import unicodedata
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from codex.config import get_settings
|
||||
from codex.models import FormulaChunk
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Unicode categories that indicate mathematical symbols.
|
||||
_MATH_CATEGORIES: frozenset[str] = frozenset({"Sm", "So"})
|
||||
|
||||
# Bbox size thresholds — avoids cropping tiny index fragments.
|
||||
_MIN_HEIGHT_PT = 15.0
|
||||
_MIN_WIDTH_PT = 40.0
|
||||
# Minimum number of math-category characters in a text block before we
|
||||
# consider it a formula candidate.
|
||||
_MIN_MATH_SYMBOLS = 5
|
||||
# Minimum ratio of math/total chars (as float 0–1).
|
||||
_MIN_MATH_RATIO = 0.15
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _is_math_char(ch: str) -> bool:
|
||||
"""Return True if *ch* belongs to a mathematical Unicode category."""
|
||||
return unicodedata.category(ch) in _MATH_CATEGORIES
|
||||
|
||||
|
||||
def _math_ratio(text: str) -> float:
|
||||
"""Return the fraction of characters in *text* that are math chars."""
|
||||
if not text:
|
||||
return 0.0
|
||||
math_count = sum(1 for ch in text if _is_math_char(ch))
|
||||
return math_count / len(text)
|
||||
|
||||
|
||||
def _math_symbol_count(text: str) -> int:
|
||||
"""Return the absolute count of math-category characters in *text*."""
|
||||
return sum(1 for ch in text if _is_math_char(ch))
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# pix2tex singleton
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_pix2tex_model: Any = None
|
||||
|
||||
|
||||
def _get_pix2tex_model() -> Any:
|
||||
"""Lazily load and return the pix2tex LatexOCR singleton."""
|
||||
global _pix2tex_model # noqa: PLW0603
|
||||
if _pix2tex_model is None:
|
||||
from pix2tex.cli import LatexOCR
|
||||
|
||||
_pix2tex_model = LatexOCR()
|
||||
return _pix2tex_model
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# pix2tex backend
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _extract_formulas_pix2tex(pdf_path: str) -> list[FormulaChunk]:
|
||||
"""Extract LaTeX formulas from *pdf_path* using pix2tex (local OCR).
|
||||
|
||||
Algorithm
|
||||
---------
|
||||
1. Open the PDF with PyMuPDF (fitz).
|
||||
2. For each page, gather text blocks; keep only blocks whose
|
||||
math-symbol ratio exceeds the threshold and whose bounding box
|
||||
is larger than the minimum size.
|
||||
3. Crop the region at 3× resolution (better OCR quality).
|
||||
4. Run pix2tex ``LatexOCR`` on the cropped image.
|
||||
5. Wrap the result in a :class:`~codex.models.FormulaChunk`.
|
||||
"""
|
||||
import fitz
|
||||
from PIL import Image
|
||||
|
||||
results: list[FormulaChunk] = []
|
||||
paper_id = Path(pdf_path).stem
|
||||
model = _get_pix2tex_model()
|
||||
|
||||
try:
|
||||
doc: Any = fitz.open(pdf_path)
|
||||
except Exception:
|
||||
logger.exception("fitz.open failed for %s", pdf_path)
|
||||
return results
|
||||
|
||||
scale = fitz.Matrix(3, 3)
|
||||
|
||||
for page_num in range(len(doc)):
|
||||
page = doc[page_num]
|
||||
blocks: list[Any] = page.get_text("blocks")
|
||||
|
||||
for block in blocks:
|
||||
# block = (x0, y0, x1, y1, text, block_no, block_type)
|
||||
if len(block) < 5:
|
||||
continue
|
||||
x0, y0, x1, y1 = block[0], block[1], block[2], block[3]
|
||||
text: str = block[4]
|
||||
|
||||
# Size filter
|
||||
width_pt = x1 - x0
|
||||
height_pt = y1 - y0
|
||||
if width_pt < _MIN_WIDTH_PT or height_pt < _MIN_HEIGHT_PT:
|
||||
continue
|
||||
|
||||
# Math content filter
|
||||
if _math_symbol_count(text) < _MIN_MATH_SYMBOLS or _math_ratio(text) < _MIN_MATH_RATIO:
|
||||
continue
|
||||
|
||||
# Crop at 3× for OCR quality
|
||||
clip_rect = fitz.Rect(x0, y0, x1, y1)
|
||||
crop_pix = page.get_pixmap(matrix=scale, clip=clip_rect)
|
||||
img = Image.frombytes(
|
||||
"RGB",
|
||||
(crop_pix.width, crop_pix.height),
|
||||
crop_pix.samples,
|
||||
)
|
||||
|
||||
try:
|
||||
latex: str = model(img)
|
||||
except Exception:
|
||||
logger.debug("pix2tex failed on page %d block, skipping", page_num + 1)
|
||||
continue
|
||||
|
||||
if not latex or not latex.strip():
|
||||
continue
|
||||
|
||||
context = text[:200].replace("\n", " ")
|
||||
results.append(
|
||||
FormulaChunk(
|
||||
paper_id=paper_id,
|
||||
page=page_num + 1,
|
||||
raw_latex=latex.strip(),
|
||||
context=context,
|
||||
)
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# MathPix backend
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _extract_formulas_mathpix(
|
||||
pdf_path: str,
|
||||
app_id: str,
|
||||
app_key: str,
|
||||
) -> list[FormulaChunk]:
|
||||
"""Extract LaTeX formulas using the MathPix v3 PDF API.
|
||||
|
||||
The function uploads the PDF, polls for results and parses the
|
||||
returned ``latex_simplified`` list into :class:`~codex.models.FormulaChunk`
|
||||
objects.
|
||||
|
||||
Retries with exponential back-off (via ``tenacity``) on transient HTTP
|
||||
errors.
|
||||
"""
|
||||
import httpx
|
||||
from tenacity import retry, stop_after_attempt, wait_exponential
|
||||
|
||||
paper_id = Path(pdf_path).stem
|
||||
headers = {
|
||||
"app_id": app_id,
|
||||
"app_key": app_key,
|
||||
}
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(4),
|
||||
wait=wait_exponential(multiplier=1, min=2, max=30),
|
||||
reraise=True,
|
||||
)
|
||||
def _post_pdf() -> dict[str, Any]:
|
||||
with open(pdf_path, "rb") as fh:
|
||||
resp = httpx.post(
|
||||
"https://api.mathpix.com/v3/pdf",
|
||||
headers=headers,
|
||||
files={"file": fh},
|
||||
timeout=60.0,
|
||||
)
|
||||
resp.raise_for_status()
|
||||
return resp.json() # type: ignore[no-any-return]
|
||||
|
||||
submit_response = _post_pdf()
|
||||
pdf_id: str = submit_response.get("pdf_id", "")
|
||||
if not pdf_id:
|
||||
logger.warning("MathPix did not return a pdf_id — aborting")
|
||||
return []
|
||||
|
||||
# Poll until complete
|
||||
import time
|
||||
|
||||
for _ in range(60):
|
||||
poll_resp = httpx.get(
|
||||
f"https://api.mathpix.com/v3/pdf/{pdf_id}.mmd",
|
||||
headers=headers,
|
||||
timeout=30.0,
|
||||
)
|
||||
if poll_resp.status_code == 200:
|
||||
break
|
||||
time.sleep(2)
|
||||
else:
|
||||
logger.warning("MathPix polling timed out for pdf_id=%s", pdf_id)
|
||||
return []
|
||||
|
||||
# Parse response — MathPix MMD contains \[ ... \] or $...$ blocks
|
||||
mmd_text: str = poll_resp.text
|
||||
results: list[FormulaChunk] = []
|
||||
lines = mmd_text.splitlines()
|
||||
in_block = False
|
||||
block_lines: list[str] = []
|
||||
page = 1
|
||||
|
||||
for line in lines:
|
||||
if line.strip().startswith("\\["):
|
||||
in_block = True
|
||||
block_lines = [line]
|
||||
elif in_block:
|
||||
block_lines.append(line)
|
||||
if line.strip().endswith("\\]"):
|
||||
latex = "\n".join(block_lines).strip()
|
||||
results.append(
|
||||
FormulaChunk(
|
||||
paper_id=paper_id,
|
||||
page=page,
|
||||
raw_latex=latex,
|
||||
context="",
|
||||
)
|
||||
)
|
||||
in_block = False
|
||||
block_lines = []
|
||||
# crude page tracking via MMD page markers
|
||||
if line.strip().startswith("<!-- Page"):
|
||||
page += 1
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Public entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def extract_formulas(
|
||||
pdf_path: str,
|
||||
mathpix_app_id: str | None = None,
|
||||
mathpix_app_key: str | None = None,
|
||||
) -> list[FormulaChunk]:
|
||||
"""Extract mathematical formulas from *pdf_path*.
|
||||
|
||||
Routes to the MathPix cloud backend when ``mathpix_app_id`` and
|
||||
``mathpix_app_key`` are both non-empty; otherwise falls back to
|
||||
the local pix2tex model (if ``PIX2TEX_FALLBACK=true``, which is
|
||||
the default).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
pdf_path:
|
||||
Absolute or relative path to the PDF file.
|
||||
mathpix_app_id:
|
||||
Override for ``MATHPIX_APP_ID`` setting (optional).
|
||||
mathpix_app_key:
|
||||
Override for ``MATHPIX_APP_KEY`` setting (optional).
|
||||
|
||||
Returns
|
||||
-------
|
||||
list[FormulaChunk]
|
||||
Extracted formulas; empty list if extraction is disabled or fails.
|
||||
"""
|
||||
settings = get_settings()
|
||||
|
||||
app_id = mathpix_app_id or settings.mathpix_app_id or ""
|
||||
app_key = mathpix_app_key or settings.mathpix_app_key or ""
|
||||
|
||||
if app_id and app_key:
|
||||
logger.info("Using MathPix backend for %s", pdf_path)
|
||||
return _extract_formulas_mathpix(pdf_path, app_id, app_key)
|
||||
|
||||
if settings.pix2tex_fallback:
|
||||
logger.info("Using pix2tex fallback backend for %s", pdf_path)
|
||||
return _extract_formulas_pix2tex(pdf_path)
|
||||
|
||||
logger.info("Formula extraction disabled (no creds, pix2tex_fallback=False)")
|
||||
return []
|
||||
Reference in New Issue
Block a user