Spaces:
Sleeping
Sleeping
File size: 3,087 Bytes
46fe385 ff72077 46fe385 ff72077 46fe385 ff72077 46fe385 ccbf2d4 46fe385 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 | """Thin wrapper around marker-pdf.
The model dict is built once at import time so subsequent requests reuse the
same weights on GPU. Don't construct PdfConverter at module level — its config
depends on the per-request `mode` (LLM enhancement on or off).
"""
from __future__ import annotations
import logging
import os
from pathlib import Path
from typing import Literal
import torch
from marker.converters.pdf import PdfConverter
from marker.models import create_model_dict
from marker.output import text_from_rendered
log = logging.getLogger(__name__)
Mode = Literal["fast", "quality"]
# Loaded once. ~5 GB on disk, lives in /data/hf in production.
_models = create_model_dict()
log.info("marker model dict loaded (cuda=%s)", torch.cuda.is_available())
CLAUDE_MODEL = os.getenv("CLAUDE_MODEL", "claude-sonnet-4-6")
def device() -> str:
return "gpu" if torch.cuda.is_available() else "cpu"
def parse(pdf_path: str | Path, mode: Mode = "fast") -> tuple[str, int]:
"""Run marker on the PDF at `pdf_path`. Returns (markdown, page_count)."""
config: dict[str, object] = {"output_format": "markdown"}
llm_service: str | None = None
if mode == "quality":
api_key = os.getenv("ANTHROPIC_API_KEY")
if not api_key:
raise RuntimeError(
"parse_mode=quality requires ANTHROPIC_API_KEY to be set."
)
config.update(
{
"use_llm": True,
"claude_api_key": api_key,
"claude_model_name": CLAUDE_MODEL,
}
)
# Marker's PdfConverter takes a dotted class path STRING here and
# resolves it via strings_to_classes(). It then constructs the
# service using the matching keys in `config` (claude_api_key etc.).
llm_service = "marker.services.claude.ClaudeService"
converter = PdfConverter(
artifact_dict=_models,
config=config,
llm_service=llm_service,
)
rendered = converter(str(pdf_path))
text, _, _ = text_from_rendered(rendered)
return text, _page_count(rendered, pdf_path)
def _page_count(rendered: object, pdf_path: str | Path) -> int:
"""Pull page count from marker's rendered metadata, fall back to pypdf.
marker 1.x exposes per-page entries on `rendered.metadata` (key has shifted
across versions — `page_stats` in recent releases). Probe a few likely
keys, then fall back to reading the source PDF.
"""
metadata = getattr(rendered, "metadata", None) or {}
if isinstance(metadata, dict):
for key in ("page_stats", "pages", "page_metadata"):
value = metadata.get(key)
if isinstance(value, list) and value:
return len(value)
if isinstance(value, int) and value > 0:
return value
try:
from pypdf import PdfReader
return len(PdfReader(str(pdf_path)).pages)
except Exception: # noqa: BLE001 — page count is best-effort metadata
log.warning("could not determine page_count for %s", pdf_path)
return 0
|