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
65
codex/cli.py
65
codex/cli.py
@@ -12,28 +12,54 @@ app = typer.Typer(help="codex — personal knowledge base for scientific papers.
|
||||
discover_app = typer.Typer(help="Discovery queries over the citation graph.")
|
||||
prov_app = typer.Typer(help="Provenance: @cite scan, code_links, bib export.")
|
||||
wiki_app = typer.Typer(help="Wiki-compile: grounded concept pages over the RAG substrate.")
|
||||
# F-09: search sub-app with paper (semantic) and formula (FTS) subcommands
|
||||
search_app = typer.Typer(
|
||||
help="Search: semantic paper search and formula full-text search.",
|
||||
invoke_without_command=True,
|
||||
)
|
||||
app.add_typer(discover_app, name="discover")
|
||||
app.add_typer(prov_app, name="provenance")
|
||||
app.add_typer(wiki_app, name="wiki")
|
||||
app.add_typer(search_app, name="search")
|
||||
|
||||
|
||||
@app.command()
|
||||
def ingest(
|
||||
paper_id: str = typer.Argument(..., help="arXiv ID, DOI, or OpenAlex W-ID"),
|
||||
source: Optional[str] = typer.Option(None, "--source", "-s", help="Path to .tex or .pdf"), # noqa: UP045
|
||||
rich: bool = typer.Option( # noqa: A002
|
||||
False,
|
||||
"--rich",
|
||||
help="Extract formulas and figures (requires PDF source, F-09).",
|
||||
),
|
||||
) -> None:
|
||||
"""Ingest a paper into the knowledge base."""
|
||||
from codex.ingest import ingest_paper
|
||||
|
||||
result = ingest_paper(paper_id, source_path=source)
|
||||
result = ingest_paper(paper_id, source_path=source, rich=rich)
|
||||
typer.echo(
|
||||
f"Ingested {result.paper_id}: {result.chunks_upserted} chunks, "
|
||||
f"{result.citations_upserted} citations"
|
||||
f"{result.citations_upserted} citations, "
|
||||
f"{result.formulas_upserted} formulas, "
|
||||
f"{result.figures_upserted} figures"
|
||||
)
|
||||
|
||||
|
||||
@app.command()
|
||||
def search(
|
||||
# ---------------------------------------------------------------------------
|
||||
# F-09: Search command group
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@search_app.callback(invoke_without_command=True)
|
||||
def search_callback(ctx: typer.Context) -> None:
|
||||
"""Search subcommands: ``paper`` for semantic search, ``formula`` for LaTeX FTS."""
|
||||
if ctx.invoked_subcommand is None:
|
||||
typer.echo(ctx.get_help())
|
||||
raise typer.Exit(0)
|
||||
|
||||
|
||||
@search_app.command("paper")
|
||||
def search_paper(
|
||||
query: str = typer.Argument(..., help="Natural-language search query"),
|
||||
limit: int = typer.Option(10, "--limit", "-n", help="Number of results"),
|
||||
) -> None:
|
||||
@@ -58,6 +84,37 @@ def search(
|
||||
typer.echo(f"[{row['distance']:.3f}] {row['id']} ({row['year']}) — {row['title']}")
|
||||
|
||||
|
||||
@search_app.command("formula")
|
||||
def search_formula(
|
||||
query: str = typer.Argument(..., help="LaTeX snippet or keyword to search in formulas"),
|
||||
limit: int = typer.Option(10, "--limit", "-n", help="Number of results"),
|
||||
) -> None:
|
||||
"""Full-text search over extracted LaTeX formulas (F-09)."""
|
||||
from codex.db import get_conn
|
||||
|
||||
with get_conn() as conn:
|
||||
rows = conn.execute(
|
||||
"""
|
||||
SELECT paper_id, page, raw_latex, context,
|
||||
ts_rank(to_tsvector('english', raw_latex || ' ' || coalesce(context, '')),
|
||||
plainto_tsquery('english', %(query)s)) AS rank
|
||||
FROM formulas
|
||||
WHERE to_tsvector('english', raw_latex || ' ' || coalesce(context, ''))
|
||||
@@ plainto_tsquery('english', %(query)s)
|
||||
ORDER BY rank DESC
|
||||
LIMIT %(limit)s
|
||||
""",
|
||||
{"query": query, "limit": limit},
|
||||
).fetchall()
|
||||
if not rows:
|
||||
typer.echo("No matching formulas found.")
|
||||
return
|
||||
for row in rows:
|
||||
typer.echo(
|
||||
f"[{row['rank']:.3f}] {row['paper_id']} p.{row['page']} {row['raw_latex'][:80]}"
|
||||
)
|
||||
|
||||
|
||||
@discover_app.command("leads")
|
||||
def discover_leads(
|
||||
limit: int = typer.Option(20, "--limit", "-n"),
|
||||
|
||||
@@ -123,6 +123,41 @@ class Settings(BaseSettings):
|
||||
),
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# F-09 Rich Parsing
|
||||
# ------------------------------------------------------------------
|
||||
mathpix_app_id: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"MathPix App ID for cloud formula extraction. "
|
||||
"When None, pix2tex local OCR is used as fallback."
|
||||
),
|
||||
)
|
||||
|
||||
mathpix_app_key: str | None = Field(
|
||||
default=None,
|
||||
description=(
|
||||
"MathPix App Key for cloud formula extraction. "
|
||||
"Must be set together with MATHPIX_APP_ID."
|
||||
),
|
||||
)
|
||||
|
||||
pix2tex_fallback: bool = Field(
|
||||
default=True,
|
||||
description=(
|
||||
"Enable pix2tex (local LaTeX OCR) as fallback when MathPix credentials "
|
||||
"are absent. Set to False to disable formula extraction entirely without creds."
|
||||
),
|
||||
)
|
||||
|
||||
figures_dir: str = Field(
|
||||
default="figures/",
|
||||
description=(
|
||||
"Directory where extracted figure images are written. "
|
||||
"Relative paths are resolved from the current working directory."
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_settings() -> Settings:
|
||||
|
||||
@@ -9,6 +9,7 @@ from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
from codex.config import get_settings
|
||||
from codex.db import get_conn
|
||||
from codex.embed import get_embedder
|
||||
from codex.models import Citation, Paper
|
||||
@@ -22,11 +23,14 @@ class IngestResult:
|
||||
paper_id: str
|
||||
chunks_upserted: int
|
||||
citations_upserted: int
|
||||
formulas_upserted: int = 0
|
||||
figures_upserted: int = 0
|
||||
|
||||
|
||||
def ingest_paper(
|
||||
paper_id: str,
|
||||
source_path: str | None = None,
|
||||
rich: bool = False,
|
||||
) -> IngestResult:
|
||||
"""Idempotent ingest of one paper.
|
||||
|
||||
@@ -42,11 +46,15 @@ def ingest_paper(
|
||||
- ``.tex`` → :func:`codex.parsing.tex.latex_to_text` + chunk
|
||||
- ``.pdf`` → :func:`codex.parsing.nougat.pdf_to_markdown` + chunk
|
||||
+ :func:`codex.parsing.grobid.extract_references` for refs
|
||||
rich:
|
||||
When True and *source_path* is a PDF, also extract formulas via
|
||||
:func:`codex.parsing.mathpix.extract_formulas` and figures via
|
||||
:func:`codex.parsing.figures.extract_figures` (F-09).
|
||||
|
||||
Returns
|
||||
-------
|
||||
IngestResult
|
||||
Counts of upserted chunks and citations.
|
||||
Counts of upserted chunks, citations, formulas, and figures.
|
||||
"""
|
||||
# ---------------------------------------------------------------
|
||||
# 1. Fetch metadata (OpenAlex primary, SemanticScholar fallback)
|
||||
@@ -195,8 +203,59 @@ def ingest_paper(
|
||||
citations_upserted = len(merged_citations)
|
||||
conn.commit()
|
||||
|
||||
# ---------------------------------------------------------------
|
||||
# 6. F-09 Rich Parsing: formulas + figures (PDF only)
|
||||
# ---------------------------------------------------------------
|
||||
formulas_upserted = 0
|
||||
figures_upserted = 0
|
||||
|
||||
if rich and source_path is not None and Path(source_path).suffix.lower() == ".pdf":
|
||||
from codex.parsing.figures import extract_figures
|
||||
from codex.parsing.mathpix import extract_formulas
|
||||
|
||||
settings = get_settings()
|
||||
|
||||
formulas = extract_formulas(source_path)
|
||||
figures = extract_figures(source_path, output_dir=settings.figures_dir)
|
||||
|
||||
if formulas or figures:
|
||||
with get_conn() as conn2:
|
||||
if formulas:
|
||||
with conn2.cursor() as cur:
|
||||
cur.executemany(
|
||||
"""
|
||||
INSERT INTO formulas (paper_id, page, raw_latex, context, eq_label)
|
||||
VALUES (%s, %s, %s, %s, %s)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
[
|
||||
(f.paper_id, f.page, f.raw_latex, f.context, f.eq_label)
|
||||
for f in formulas
|
||||
],
|
||||
)
|
||||
formulas_upserted = len(formulas)
|
||||
|
||||
if figures:
|
||||
with conn2.cursor() as cur:
|
||||
cur.executemany(
|
||||
"""
|
||||
INSERT INTO figures (paper_id, page, caption, image_path)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
ON CONFLICT DO NOTHING
|
||||
""",
|
||||
[
|
||||
(fig.paper_id, fig.page, fig.caption, fig.image_path)
|
||||
for fig in figures
|
||||
],
|
||||
)
|
||||
figures_upserted = len(figures)
|
||||
|
||||
conn2.commit()
|
||||
|
||||
return IngestResult(
|
||||
paper_id=paper.id,
|
||||
chunks_upserted=chunks_upserted,
|
||||
citations_upserted=citations_upserted,
|
||||
formulas_upserted=formulas_upserted,
|
||||
figures_upserted=figures_upserted,
|
||||
)
|
||||
|
||||
@@ -78,3 +78,39 @@ class CodeLink:
|
||||
note: str | None = None
|
||||
id: int | None = None
|
||||
added_at: datetime | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class FormulaChunk:
|
||||
"""Maps to the ``formulas`` table (F-09 Rich Parsing).
|
||||
|
||||
Stores a single extracted mathematical formula (LaTeX) from a PDF page.
|
||||
``id`` is set by the database (BIGSERIAL).
|
||||
``embedding`` is reserved for future pgvector similarity search.
|
||||
"""
|
||||
|
||||
paper_id: str
|
||||
page: int
|
||||
raw_latex: str
|
||||
context: str
|
||||
id: int | None = None
|
||||
eq_label: str | None = None
|
||||
embedding: list[float] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class FigureChunk:
|
||||
"""Maps to the ``figures`` table (F-09 Rich Parsing).
|
||||
|
||||
Stores metadata for a single extracted figure from a PDF page.
|
||||
``image_path`` points to the saved PNG on disk.
|
||||
``id`` is set by the database (BIGSERIAL).
|
||||
``embedding`` is reserved for future pgvector similarity search.
|
||||
"""
|
||||
|
||||
paper_id: str
|
||||
page: int
|
||||
image_path: str
|
||||
caption: str
|
||||
id: int | None = None
|
||||
embedding: list[float] | None = None
|
||||
|
||||
166
codex/parsing/figures.py
Normal file
166
codex/parsing/figures.py
Normal file
@@ -0,0 +1,166 @@
|
||||
"""Figure extraction from PDF pages.
|
||||
|
||||
Uses PyMuPDF (``fitz``) to iterate over embedded images in each page,
|
||||
render them to PNG and detect nearby captions using text-block heuristics.
|
||||
|
||||
The public entry point is :func:`extract_figures`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from codex.models import FigureChunk
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Caption prefixes (case-insensitive check against text block content).
|
||||
_CAPTION_PREFIXES: tuple[str, ...] = ("figure", "fig.", "abbildung")
|
||||
|
||||
# Maximum vertical distance (in points) from the image bottom or top
|
||||
# to consider a text block as the caption.
|
||||
_CAPTION_TOLERANCE_Y = 60.0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Internal helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _looks_like_caption(text: str) -> bool:
|
||||
"""Return True if *text* starts with a known caption prefix."""
|
||||
lowered = text.strip().lower()
|
||||
return any(lowered.startswith(prefix) for prefix in _CAPTION_PREFIXES)
|
||||
|
||||
|
||||
def _find_caption(
|
||||
img_rect: Any, # fitz.Rect
|
||||
text_blocks: list[Any],
|
||||
) -> str:
|
||||
"""Find the closest caption text block to *img_rect*.
|
||||
|
||||
Searches text blocks that:
|
||||
1. Start with a caption prefix (e.g. "Figure", "Fig.", "Abbildung").
|
||||
2. Are within ``_CAPTION_TOLERANCE_Y`` points above or below the image.
|
||||
|
||||
Returns the first matching text, stripped; empty string if none found.
|
||||
"""
|
||||
best: str = ""
|
||||
best_dist = float("inf")
|
||||
|
||||
for block in text_blocks:
|
||||
if len(block) < 5:
|
||||
continue
|
||||
bx0, by0, bx1, by1 = block[0], block[1], block[2], block[3]
|
||||
text: str = block[4]
|
||||
|
||||
if not _looks_like_caption(text):
|
||||
continue
|
||||
|
||||
# Distance: gap below image or gap above image
|
||||
gap_below = by0 - img_rect.y1
|
||||
gap_above = img_rect.y0 - by1
|
||||
|
||||
if 0 <= gap_below <= _CAPTION_TOLERANCE_Y:
|
||||
dist = gap_below
|
||||
elif 0 <= gap_above <= _CAPTION_TOLERANCE_Y:
|
||||
dist = gap_above
|
||||
else:
|
||||
continue
|
||||
|
||||
# Horizontal overlap sanity check (caption should overlap with image)
|
||||
overlap = min(bx1, img_rect.x1) - max(bx0, img_rect.x0)
|
||||
if overlap < 0:
|
||||
continue
|
||||
|
||||
if dist < best_dist:
|
||||
best_dist = dist
|
||||
best = text.replace("\n", " ").strip()
|
||||
|
||||
return best
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Public entry point
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def extract_figures(
|
||||
pdf_path: str,
|
||||
output_dir: str,
|
||||
) -> list[FigureChunk]:
|
||||
"""Extract embedded figures from *pdf_path* and save them as PNG files.
|
||||
|
||||
For each image on each page:
|
||||
1. Extract the image via ``page.get_images(full=True)``.
|
||||
2. Render the image as a PNG via ``fitz.Pixmap``.
|
||||
3. Search nearby text blocks for a caption (``Figure …`` / ``Fig. …`` /
|
||||
``Abbildung …``).
|
||||
4. Save the PNG to ``<output_dir>/<paper_id>_p<page>_fig<n>.png``.
|
||||
5. Yield a :class:`~codex.models.FigureChunk`.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
pdf_path:
|
||||
Path to the source PDF.
|
||||
output_dir:
|
||||
Directory where PNG images are written (created if absent).
|
||||
|
||||
Returns
|
||||
-------
|
||||
list[FigureChunk]
|
||||
Extracted figure metadata; empty list if no images found.
|
||||
"""
|
||||
import fitz
|
||||
|
||||
results: list[FigureChunk] = []
|
||||
paper_id = Path(pdf_path).stem
|
||||
out_dir = Path(output_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
try:
|
||||
doc: Any = fitz.open(pdf_path)
|
||||
except Exception:
|
||||
logger.exception("fitz.open failed for %s", pdf_path)
|
||||
return results
|
||||
|
||||
for page_num in range(len(doc)):
|
||||
page = doc[page_num]
|
||||
images: list[Any] = page.get_images(full=True)
|
||||
text_blocks: list[Any] = page.get_text("blocks")
|
||||
|
||||
for fig_idx, img_info in enumerate(images):
|
||||
xref: int = img_info[0]
|
||||
|
||||
# Get image bounding box on the page
|
||||
img_rects: list[Any] = page.get_image_rects(xref)
|
||||
if not img_rects:
|
||||
continue
|
||||
img_rect = img_rects[0]
|
||||
|
||||
# Extract and save the image pixmap
|
||||
try:
|
||||
pix: Any = fitz.Pixmap(doc, xref)
|
||||
if pix.n > 4: # CMYK or similar — convert to RGB
|
||||
pix = fitz.Pixmap(fitz.csRGB, pix)
|
||||
filename = f"{paper_id}_p{page_num + 1}_fig{fig_idx + 1}.png"
|
||||
img_path = out_dir / filename
|
||||
pix.save(str(img_path))
|
||||
except Exception:
|
||||
logger.debug("Could not extract image xref=%d on page %d", xref, page_num + 1)
|
||||
continue
|
||||
|
||||
caption = _find_caption(img_rect, text_blocks)
|
||||
|
||||
results.append(
|
||||
FigureChunk(
|
||||
paper_id=paper_id,
|
||||
page=page_num + 1,
|
||||
image_path=str(img_path),
|
||||
caption=caption,
|
||||
)
|
||||
)
|
||||
|
||||
return results
|
||||
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