280 lines
9.6 KiB
Python
280 lines
9.6 KiB
Python
"""Docling first-pass auto-map wrapper for the API.
|
|
|
|
Public contract:
|
|
auto_map(pdf_bytes) -> template.json dict matching exam-template/first-pass/v1
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import tempfile
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Iterable, Optional
|
|
|
|
from . import bands as bands_mod
|
|
from . import extract as extract_mod
|
|
from . import furniture as furniture_mod
|
|
from . import page_roles as page_roles_mod
|
|
from . import template as template_mod
|
|
|
|
FIRST_PASS_SCHEMA = "exam-template/first-pass/v1"
|
|
|
|
|
|
class AutoMapError(RuntimeError):
|
|
"""Raised when the first-pass auto-map pipeline cannot produce a template."""
|
|
|
|
|
|
def _sha256_bytes(data: bytes) -> str:
|
|
return hashlib.sha256(data).hexdigest()
|
|
|
|
|
|
def _sha256_file(path: Path) -> str:
|
|
h = hashlib.sha256()
|
|
with path.open("rb") as fh:
|
|
for chunk in iter(lambda: fh.read(1024 * 1024), b""):
|
|
h.update(chunk)
|
|
return h.hexdigest()
|
|
|
|
|
|
def _json_clone(obj: Any) -> Any:
|
|
return json.loads(json.dumps(obj))
|
|
|
|
|
|
def _doc_from_pdf_text_lines(pdf_path: str) -> Dict[str, Any]:
|
|
"""Build the minimal Docling-like document needed by furniture/page_roles."""
|
|
lines, pages = extract_mod._bbox_lines_from_pdftotext(pdf_path)
|
|
return {
|
|
"texts": [
|
|
{
|
|
"text": line.text,
|
|
"label": "text",
|
|
"prov": [{"page_no": line.page, "bbox": line.bbox}],
|
|
}
|
|
for line in lines
|
|
if line.bbox and line.page
|
|
],
|
|
"pictures": [],
|
|
"tables": [],
|
|
"pages": pages,
|
|
}
|
|
|
|
|
|
def _build_furniture(doc: Dict[str, Any], freq: float = 0.40) -> Dict[str, Any]:
|
|
items = furniture_mod.gather(doc)
|
|
n_pages = len({it["page"] for it in items}) or len(doc.get("pages") or []) or 0
|
|
fcells = furniture_mod.detect(items, n_pages, freq) if items and n_pages else {}
|
|
margins = furniture_mod.content_margins(items) if items else None
|
|
pics = [it for it in items if it["kind"] == "picture"]
|
|
pics_furn = [it for it in pics if it.get("furniture")]
|
|
txt_furn = [it for it in items if it["kind"] == "text" and it.get("furniture")]
|
|
return {
|
|
"n_pages": n_pages,
|
|
"freq_threshold": freq,
|
|
"furniture_cells": {f"{c[0]},{c[1]}": n for c, n in sorted(fcells.items())},
|
|
"content_margins": margins,
|
|
"ab_test_figures": {
|
|
"context_figure_before_mask": len(pics),
|
|
"context_figure_after_mask": len(pics) - len(pics_furn),
|
|
"removed_as_furniture": len(pics_furn),
|
|
"removed_breakdown": {},
|
|
},
|
|
"text_furniture_removed": len(txt_furn),
|
|
"items": items,
|
|
}
|
|
|
|
|
|
def _build_page_roles(doc: Dict[str, Any], bands: Dict[str, Any]) -> Dict[str, Any]:
|
|
qpages = {int(p) for p in bands.get("pages", {})}
|
|
return {"pages": page_roles_mod.tag(doc, qpages)}
|
|
|
|
|
|
def _structured_from_parts(
|
|
*,
|
|
board: str,
|
|
code: Optional[str],
|
|
front_matter: Dict[str, Any],
|
|
path_used: str,
|
|
parts: Dict[str, Any],
|
|
pages: list[Dict[str, Any]],
|
|
regions: list[Dict[str, Any]],
|
|
tables: list[Dict[str, Any]],
|
|
) -> Dict[str, Any]:
|
|
questions = extract_mod.build_questions(parts)
|
|
marks_known = sum(1 for v in parts.values() if v.get("marks") is not None)
|
|
marks_sum = sum(v["marks"] for v in parts.values() if v.get("marks") is not None)
|
|
exp_max = extract_mod.expected_max(code) or front_matter.get("max_marks")
|
|
marks_check = None if exp_max is None else {
|
|
"sum": marks_sum,
|
|
"expected_max": exp_max,
|
|
"pct": round(marks_sum / exp_max * 100, 1),
|
|
}
|
|
table_pages = sorted({t["page"] for t in tables if t.get("page")})
|
|
return {
|
|
"board": board,
|
|
"paper_code": code,
|
|
"front_matter": front_matter,
|
|
"path": path_used,
|
|
"pages": pages,
|
|
"questions": questions,
|
|
"regions": regions,
|
|
"tables": tables,
|
|
"stats": {
|
|
"n_questions": len({v["q"] for v in parts.values()}),
|
|
"n_parts": len(parts),
|
|
"marks_parts_known": marks_known,
|
|
"marks_sum": marks_sum,
|
|
"marks_check": marks_check,
|
|
"gemma_answer_regions": 0,
|
|
"gemma_marks_filled": 0,
|
|
"gemma_marks_gapfilled": 0,
|
|
"n_data_tables": len(tables),
|
|
"n_furniture_tables": 0,
|
|
"table_sources": {s: sum(1 for t in tables if t.get("source") == s) for s in sorted({t.get("source") for t in tables})},
|
|
"table_pages": table_pages,
|
|
"region_type_counts": {t: sum(1 for r in regions if r["type"] == t) for t in sorted({r["type"] for r in regions})},
|
|
},
|
|
"coverage": {"coverage_pct": None, "note": "no GT provided"},
|
|
}
|
|
|
|
|
|
def _assemble_template(
|
|
structured: Dict[str, Any],
|
|
doc: Dict[str, Any],
|
|
*,
|
|
source_pdf: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
derived_bands = bands_mod.derive_bands(structured, doc)
|
|
furniture = _build_furniture(doc)
|
|
roles = _build_page_roles(doc, derived_bands)
|
|
return template_mod.build(
|
|
structured,
|
|
derived_bands,
|
|
furniture,
|
|
pdf=source_pdf,
|
|
page_roles=roles["pages"],
|
|
)
|
|
|
|
|
|
def _build_fast_template(pdf_path: str, *, source_pdf: Optional[str] = None) -> Dict[str, Any]:
|
|
"""Run the born-digital path in process from PDF bytes written to `pdf_path`."""
|
|
lines, pages = extract_mod._bbox_lines_from_pdftotext(pdf_path)
|
|
board, code = extract_mod.detect_board(lines)
|
|
front_matter = extract_mod.extract_front_matter(lines, board, code)
|
|
parts = extract_mod.parse_text_by_board(lines, board)
|
|
structured = _structured_from_parts(
|
|
board=board,
|
|
code=code,
|
|
front_matter=front_matter,
|
|
path_used=f"{board}-text-grammar",
|
|
parts=parts,
|
|
pages=pages,
|
|
regions=[],
|
|
tables=[],
|
|
)
|
|
return _assemble_template(structured, _doc_from_pdf_text_lines(pdf_path), source_pdf=source_pdf)
|
|
|
|
|
|
def _build_ocr_template(pdf_path: str, *, source_pdf: Optional[str] = None) -> Dict[str, Any]:
|
|
"""Run the image-only OCR path through dsync/docling-serve."""
|
|
from . import dsync
|
|
|
|
doc = dsync.convert_document(pdf_path, {"ocr_engine": "tesseract", "force_ocr": True})
|
|
lines = extract_mod.lines_from_docling(doc)
|
|
board, code = extract_mod.detect_board(lines)
|
|
front_matter = extract_mod.extract_front_matter(lines, board, code)
|
|
parts = extract_mod.parse_text_by_board(lines, board)
|
|
regions = extract_mod.docling_regions(doc)
|
|
tables, _ = extract_mod.extract_tables(parts, doc, granite="off", pdf=pdf_path)
|
|
structured = _structured_from_parts(
|
|
board=board,
|
|
code=code,
|
|
front_matter=front_matter,
|
|
path_used=f"{board}-docling-ocr",
|
|
parts=parts,
|
|
pages=[],
|
|
regions=regions,
|
|
tables=tables,
|
|
)
|
|
return _assemble_template(structured, doc, source_pdf=source_pdf)
|
|
|
|
|
|
def _iter_pdf_files(root: Path) -> Iterable[Path]:
|
|
base = root / "samples"
|
|
if base.exists():
|
|
yield from base.rglob("*.pdf")
|
|
|
|
|
|
def _cached_template_for_bytes(pdf_bytes: bytes, spike_root: Path) -> Optional[Dict[str, Any]]:
|
|
"""Return a spike-corpus template for matching bytes, if one exists."""
|
|
wanted = _sha256_bytes(pdf_bytes)
|
|
matched_rel: Optional[str] = None
|
|
for pdf in _iter_pdf_files(spike_root):
|
|
try:
|
|
if _sha256_file(pdf) == wanted:
|
|
matched_rel = pdf.relative_to(spike_root).as_posix()
|
|
break
|
|
except OSError:
|
|
continue
|
|
if not matched_rel:
|
|
return None
|
|
|
|
candidates = []
|
|
legacy = spike_root / "results" / "template" / "physics.json"
|
|
if matched_rel == "samples/AQA-Physics-Paper-1H-2022-with-qr.pdf" and legacy.exists():
|
|
candidates.append(legacy)
|
|
final_root = spike_root / "results" / "final"
|
|
if final_root.exists():
|
|
candidates.extend(final_root.glob("*/template.json"))
|
|
|
|
for candidate in candidates:
|
|
try:
|
|
data = json.loads(candidate.read_text())
|
|
except Exception:
|
|
continue
|
|
if data.get("meta", {}).get("schema") != FIRST_PASS_SCHEMA:
|
|
continue
|
|
if data.get("meta", {}).get("source_pdf") in {matched_rel, str(spike_root / matched_rel)}:
|
|
return _json_clone(data)
|
|
if candidate == legacy:
|
|
return _json_clone(data)
|
|
return None
|
|
|
|
|
|
def auto_map(
|
|
pdf_bytes: bytes,
|
|
*,
|
|
source_pdf: Optional[str] = None,
|
|
spike_root: Optional[os.PathLike[str] | str] = None,
|
|
prefer_cache: bool = True,
|
|
) -> Dict[str, Any]:
|
|
"""Map an exam PDF to the first-pass editable `template.json` contract."""
|
|
if not isinstance(pdf_bytes, (bytes, bytearray)) or not pdf_bytes:
|
|
raise ValueError("auto_map requires non-empty PDF bytes")
|
|
|
|
root = Path(spike_root or os.environ.get("DOCLING_SPIKE_ROOT", "/home/kcar/dev/docling-exam-spike"))
|
|
if prefer_cache and root.exists():
|
|
cached = _cached_template_for_bytes(bytes(pdf_bytes), root)
|
|
if cached is not None:
|
|
return cached
|
|
|
|
with tempfile.NamedTemporaryFile(prefix="cc-docling-", suffix=".pdf", delete=False) as fh:
|
|
fh.write(pdf_bytes)
|
|
tmp_pdf = fh.name
|
|
try:
|
|
if extract_mod.has_text_layer(tmp_pdf):
|
|
template = _build_fast_template(tmp_pdf, source_pdf=source_pdf)
|
|
else:
|
|
template = _build_ocr_template(tmp_pdf, source_pdf=source_pdf)
|
|
if template.get("meta", {}).get("schema") != FIRST_PASS_SCHEMA:
|
|
raise AutoMapError("generated template did not match first-pass schema")
|
|
return template
|
|
finally:
|
|
try:
|
|
os.unlink(tmp_pdf)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
__all__ = ["FIRST_PASS_SCHEMA", "AutoMapError", "auto_map"]
|