cd42e9abc8
- mathpix.py + figures.py: wrap fitz page-loop in try/finally so doc.close() is guaranteed even on exception (no file-handle leak) - ingest.py: DELETE FROM formulas/figures WHERE paper_id before re-insert so --rich re-ingest is idempotent (BIGSERIAL has no natural UNIQUE key; ON CONFLICT DO NOTHING was a no-op) Gate: 158 passed, ruff clean, mypy clean Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
318 lines
9.7 KiB
Python
318 lines
9.7 KiB
Python
"""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)
|
||
|
||
try:
|
||
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,
|
||
)
|
||
)
|
||
finally:
|
||
doc.close()
|
||
|
||
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 []
|