Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| """ | |
| test_docling_ocr.py β Test OCR on PeerRead ACL-2017 test PDFs. | |
| ChαΊ‘y tα»« thΖ° mα»₯c docs/: | |
| python scripts/test_docling_ocr.py --provider docling-granite # Granite VLM qua Kaggle (khuyαΊΏn nghα») | |
| python scripts/test_docling_ocr.py --provider docling-granite-local # Granite VLM chαΊ‘y local (cαΊ§n GPU) | |
| python scripts/test_docling_ocr.py --provider local # PyMuPDF baseline | |
| python scripts/test_docling_ocr.py --compare | |
| Output Δược linearize vα» ΔΓΊng template plain-text cα»§a ground truth (test/parsed_pdfs): | |
| bαΊ£ng lΓ m phαΊ³ng (cell cΓ‘ch space, 1 dΓ²ng/row), cΓ΄ng thα»©c vα» unicode thΖ°α»ng (khΓ΄ng LaTeX), | |
| chα» giα»― heading "# "/"## " Δα» tΓ‘ch section. | |
| HuggingFace token (tΓΉy chα»n, giΓΊp download model α»n Δα»nh hΖ‘n): | |
| ΔαΊ·t HF_TOKEN=hf_... trong file .env α» root project, hoαΊ·c: | |
| set HF_TOKEN=hf_... (Windows) | |
| export HF_TOKEN=hf_... (Linux/Mac) | |
| """ | |
| from __future__ import annotations | |
| # Suppress TensorFlow/oneDNN noise β must be set before any import | |
| import os | |
| os.environ.setdefault("TF_ENABLE_ONEDNN_OPTS", "0") | |
| os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3") | |
| import argparse | |
| import asyncio | |
| import json | |
| import re | |
| import sys | |
| import unicodedata | |
| import warnings | |
| warnings.filterwarnings("ignore", category=UserWarning) | |
| warnings.filterwarnings("ignore", message=".*NumPy.*") | |
| from difflib import SequenceMatcher | |
| from pathlib import Path | |
| # Load .env tα»« root project (chα»©a HF_TOKEN, ANTHROPIC_API_KEY, ...) | |
| _root = Path(__file__).resolve().parent.parent.parent | |
| _env_file = _root / ".env" | |
| if _env_file.exists(): | |
| for _line in _env_file.read_text(encoding="utf-8").splitlines(): | |
| _line = _line.strip() | |
| if _line and not _line.startswith("#") and "=" in _line: | |
| _k, _, _v = _line.partition("=") | |
| os.environ.setdefault(_k.strip(), _v.strip()) | |
| sys.path.insert(0, str(_root)) | |
| from backend.pipeline import pdf2md as pdf2md_backend | |
| # ββ DoclingDocument β ground-truth-style plain text βββββββββββββββββββββββββββ | |
| # The ground truth in test/parsed_pdfs (GROBID/CRF) is PLAIN TEXT: | |
| # β’ tables are flattened into the section text, cells joined by spaces, one row | |
| # per line (NO markdown pipe-tables); | |
| # β’ formulas are plain unicode (e.g. "ri = Ο(Wrxi + Urhiβ1 + br)"), NOT LaTeX | |
| # and NOT Docling's "<!-- formula-not-decoded -->" placeholder; | |
| # β’ pictures / page headers / footers are dropped. | |
| # `result.document.export_to_markdown()` instead emits pipe-tables, "$$...$$" | |
| # LaTeX and image placeholders, so it does not line up with the ground truth. | |
| # We therefore linearize the DoclingDocument ourselves to the same template, | |
| # keeping only "# "/"## " headings so md_to_peerread can still split sections. | |
| _GREEK = { | |
| r"\alpha": "Ξ±", r"\beta": "Ξ²", r"\gamma": "Ξ³", r"\delta": "Ξ΄", | |
| r"\epsilon": "Ξ΅", r"\varepsilon": "Ξ΅", r"\zeta": "ΞΆ", r"\eta": "Ξ·", | |
| r"\theta": "ΞΈ", r"\iota": "ΞΉ", r"\kappa": "ΞΊ", r"\lambda": "Ξ»", | |
| r"\mu": "ΞΌ", r"\nu": "Ξ½", r"\xi": "ΞΎ", r"\pi": "Ο", r"\rho": "Ο", | |
| r"\sigma": "Ο", r"\tau": "Ο", r"\upsilon": "Ο ", r"\phi": "Ο", | |
| r"\varphi": "Ο", r"\chi": "Ο", r"\psi": "Ο", r"\omega": "Ο", | |
| r"\Gamma": "Ξ", r"\Delta": "Ξ", r"\Theta": "Ξ", r"\Lambda": "Ξ", | |
| r"\Sigma": "Ξ£", r"\Phi": "Ξ¦", r"\Psi": "Ξ¨", r"\Omega": "Ξ©", | |
| } | |
| _SYMBOLS = { | |
| r"\times": "Γ", r"\cdot": "Β·", r"\cdots": "Β·Β·Β·", r"\ldots": "...", | |
| r"\leq": "β€", r"\geq": "β₯", r"\neq": "β ", r"\approx": "β", | |
| r"\rightarrow": "β", r"\to": "β", r"\leftarrow": "β", r"\Rightarrow": "β", | |
| r"\infty": "β", r"\sum": "β", r"\prod": "β", r"\int": "β«", | |
| r"\partial": "β", r"\nabla": "β", r"\in": "β", r"\notin": "β", | |
| r"\forall": "β", r"\exists": "β", r"\pm": "Β±", r"\mp": "β", | |
| } | |
| _LATEX_MAP = {**_GREEK, **_SYMBOLS} | |
| def _latex_to_plain(latex: str) -> str: | |
| """Best-effort LaTeX β plain unicode, to match GROBID's embedded-text style. | |
| e.g. r_i = \\sigma(W_r x_i + U_r h_{i-1} + b_r) -> ri = Ο(Wr xi + Ur hi-1 + br) | |
| Not a full LaTeX renderer β just enough to make the OCR output comparable | |
| with the plain-text ground truth. | |
| """ | |
| s = (latex or "").strip() | |
| if not s: | |
| return "" | |
| s = s.replace("$$", "").replace("$", "") | |
| s = re.sub(r"\\[,;:!>\s]", " ", s) # spacing macros: \, \; \! | |
| s = s.replace(r"\left", "").replace(r"\right", "") | |
| for k in sorted(_LATEX_MAP, key=len, reverse=True): # longest first | |
| s = s.replace(k, _LATEX_MAP[k]) | |
| s = re.sub(r"\\frac\s*\{([^{}]*)\}\s*\{([^{}]*)\}", r"(\1)/(\2)", s) | |
| s = re.sub(r"\\sqrt\s*\{([^{}]*)\}", r"β(\1)", s) | |
| s = re.sub( | |
| r"\\(?:mathrm|mathbf|mathbb|mathcal|text|operatorname|textbf|textit|boldsymbol)\s*\{([^{}]*)\}", | |
| r"\1", s, | |
| ) | |
| # subscripts / superscripts: drop the _ / ^ marker, keep the content inline | |
| s = re.sub(r"[_^]\{([^{}]*)\}", r"\1", s) | |
| s = re.sub(r"[_^]\s*(\w)", r"\1", s) | |
| s = re.sub(r"\\([A-Za-z]+)", r"\1", s) # leftover commands: keep name | |
| s = s.replace("{", "").replace("}", "") | |
| return re.sub(r"\s+", " ", s).strip() | |
| def _table_to_plain(table_item) -> str: | |
| """Flatten a Docling table to GROBID style: cells space-joined, one row/line.""" | |
| try: | |
| df = table_item.export_to_dataframe() | |
| except Exception: | |
| return "" | |
| if df is None or getattr(df, "empty", True): | |
| return "" | |
| import pandas as pd | |
| # GROBID flattens the WHOLE table onto one line: header cells then every row, | |
| # all space-joined (e.g. "Corpus # words # chunks # sentences Train ... Test ..."). | |
| cells = [str(c).strip() for c in df.columns] | |
| for _, row in df.iterrows(): | |
| cells.extend("" if pd.isna(v) else str(v).strip() for v in row.tolist()) | |
| return re.sub(r"\s+", " ", " ".join(cells)).strip() | |
| def _docling_doc_to_text(doc) -> str: | |
| """Linearize a DoclingDocument to GT-style plain text (with #/## headings).""" | |
| from docling_core.types.doc import DocItemLabel, TableItem | |
| # Keep CAPTION β GROBID keeps "Table 1: ...", "Figure 2: ..." inline in the text. | |
| drop = { | |
| DocItemLabel.PAGE_HEADER, DocItemLabel.PAGE_FOOTER, DocItemLabel.PICTURE, | |
| } | |
| parts: list[str] = [] | |
| for item, _level in doc.iterate_items(): | |
| if isinstance(item, TableItem): | |
| tbl = _table_to_plain(item) | |
| if tbl: | |
| parts.append(tbl) | |
| continue | |
| label = getattr(item, "label", None) | |
| if label in drop: | |
| continue | |
| text = (getattr(item, "text", "") or "").strip() | |
| # Never leak HTML comments such as the VLM's "<!-- formula-not-decoded -->" | |
| # placeholder β strip them so undecoded formulas simply vanish. | |
| text = re.sub(r"<!--.*?-->", "", text, flags=re.S).strip() | |
| if label == DocItemLabel.FORMULA: | |
| text = _latex_to_plain(text) | |
| if text: | |
| parts.append(text) | |
| continue | |
| if not text: | |
| continue | |
| if label == DocItemLabel.TITLE: | |
| parts.append(f"# {text}") | |
| elif label == DocItemLabel.SECTION_HEADER: | |
| parts.append(f"## {text}") | |
| else: | |
| parts.append(text) | |
| return "\n\n".join(parts) | |
| # ββ Docling local Granite VLM provider (Python API, no server) βββββββββββββββββ | |
| def _login_huggingface_if_available() -> None: | |
| """Authenticate once if HF_TOKEN is present, then continue without auth on failure.""" | |
| hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN") | |
| if not hf_token: | |
| return | |
| try: | |
| from huggingface_hub import login | |
| login(token=hf_token, add_to_git_credential=False) | |
| except Exception: | |
| pass | |
| # ββ Docling Granite VLM (full-page "Convert this page to docling." preset) βββββ | |
| def _pdf2md_docling_granite_local(pdf_bytes: bytes) -> str: | |
| """Run Docling's VLM pipeline with the Granite Docling preset.""" | |
| import tempfile | |
| from pathlib import Path as _Path | |
| _login_huggingface_if_available() | |
| from docling.datamodel.base_models import InputFormat | |
| from docling.datamodel.pipeline_options import ( | |
| AcceleratorDevice, | |
| AcceleratorOptions, | |
| VlmConvertOptions, | |
| VlmPipelineOptions, | |
| ) | |
| from docling.document_converter import DocumentConverter, PdfFormatOption | |
| from docling.pipeline.vlm_pipeline import VlmPipeline | |
| vlm_options = VlmConvertOptions.from_preset("granite_docling") | |
| pipeline_options = VlmPipelineOptions(vlm_options=vlm_options) | |
| # AUTO picks CUDA/MPS if present, else CPU (slow β local is only for spot checks). | |
| pipeline_options.accelerator_options = AcceleratorOptions(device=AcceleratorDevice.AUTO) | |
| converter = DocumentConverter( | |
| format_options={ | |
| InputFormat.PDF: PdfFormatOption( | |
| pipeline_cls=VlmPipeline, | |
| pipeline_options=pipeline_options, | |
| ) | |
| } | |
| ) | |
| with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: | |
| tmp.write(pdf_bytes) | |
| tmp_path = _Path(tmp.name) | |
| try: | |
| result = converter.convert(str(tmp_path)) | |
| return _docling_doc_to_text(result.document) | |
| finally: | |
| tmp_path.unlink(missing_ok=True) | |
| # ββ Docling standard pipeline + formula/table enrichment (decodes formulas) ββββ | |
| def _pdf2md_docling_standard_local(pdf_bytes: bytes) -> str: | |
| """Standard PDF pipeline with formula enrichment (formulas β LaTeX) and | |
| TableFormer (ACCURATE). Use this when the VLM leaves formulas undecoded β | |
| the CodeFormula model decodes them so `_latex_to_plain` can render them.""" | |
| import tempfile | |
| from pathlib import Path as _Path | |
| _login_huggingface_if_available() | |
| from docling.datamodel.base_models import InputFormat | |
| from docling.datamodel.pipeline_options import ( | |
| AcceleratorDevice, | |
| AcceleratorOptions, | |
| PdfPipelineOptions, | |
| TableFormerMode, | |
| TableStructureOptions, | |
| ) | |
| from docling.document_converter import DocumentConverter, PdfFormatOption | |
| po = PdfPipelineOptions() | |
| po.do_ocr = False # digital ACL PDFs already carry a text layer | |
| po.do_table_structure = True | |
| po.table_structure_options = TableStructureOptions(mode=TableFormerMode.ACCURATE) | |
| po.do_formula_enrichment = True # decode formulas β LaTeX (CodeFormula model) | |
| po.generate_page_images = False | |
| po.accelerator_options = AcceleratorOptions(device=AcceleratorDevice.AUTO) | |
| converter = DocumentConverter( | |
| format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=po)} | |
| ) | |
| with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp: | |
| tmp.write(pdf_bytes) | |
| tmp_path = _Path(tmp.name) | |
| try: | |
| result = converter.convert(str(tmp_path)) | |
| return _docling_doc_to_text(result.document) | |
| finally: | |
| tmp_path.unlink(missing_ok=True) | |
| async def _pdf2md_docling_server(pdf_bytes: bytes) -> str: | |
| """Send PDF to the Kaggle Docling server (Granite VLM) and get plain text back. | |
| The server linearizes the DoclingDocument to GT-style plain text (same | |
| template as test/parsed_pdfs), so the returned "markdown" is really the | |
| flattened text with only "# "/"## " headings preserved. | |
| """ | |
| import httpx | |
| async def _post(url: str) -> "httpx.Response": | |
| async with httpx.AsyncClient(timeout=600) as client: | |
| return await client.post( | |
| f"{url}/parse", | |
| files={"file": ("paper.pdf", pdf_bytes, "application/pdf")}, | |
| ) | |
| base_url = pdf2md_backend._get_docling_url() | |
| resp = await _post(base_url) | |
| # Stale tunnel URL β re-discover via ntfy and retry once. | |
| if resp.status_code in (502, 503, 404): | |
| pdf2md_backend._docling_url_cache = None | |
| pdf2md_backend._cache_fetched_at = 0.0 | |
| resp = await _post(pdf2md_backend._get_docling_url()) | |
| resp.raise_for_status() | |
| return resp.json()["markdown"] | |
| # ββ Local PDFβMarkdown provider (PyMuPDF) βββββββββββββββββββββββββββββββββββββ | |
| def _pdf2md_local(pdf_bytes: bytes) -> str: | |
| """ | |
| Convert PDF to Markdown using PyMuPDF structured text extraction. | |
| Uses font size and bold flags to detect headings vs body text. | |
| """ | |
| import fitz # PyMuPDF | |
| import statistics | |
| doc = fitz.open(stream=pdf_bytes, filetype="pdf") | |
| # ββ Pass 1: collect all font sizes to determine thresholds ββββββββββββββ | |
| all_sizes: list[float] = [] | |
| for page in doc: | |
| for block in page.get_text("dict")["blocks"]: | |
| if block.get("type") != 0: | |
| continue | |
| for line in block.get("lines", []): | |
| for span in line.get("spans", []): | |
| txt = span["text"].strip() | |
| if txt and len(txt) > 2: | |
| all_sizes.append(span["size"]) | |
| if not all_sizes: | |
| doc.close() | |
| return "" | |
| body_size = statistics.median(all_sizes) | |
| # Thresholds calibrated for ACL two-column paper style: | |
| # body ~10-11pt, section headings ~12pt (1.10x), title ~14pt (1.28x) | |
| h1_threshold = body_size * 1.25 # title / very large heading | |
| h2_threshold = body_size * 1.07 # section heading size baseline | |
| # Pattern for ACL-style section headings (content-based check). | |
| # Two tiers: | |
| # 1. Numbered headings: "1 Introduction", "2.1 Data" β very reliable. | |
| # 2. Keyword-prefix headings β keyword at start, short line (<= 60 chars). | |
| # Matches "Introduction", "Introduction and Motivation", etc. | |
| # Does NOT match long sentences that start with a keyword word. | |
| _numbered_heading_re = re.compile( | |
| r"^\d+(?:\.\d+)*\.?\s+[A-Z\d]", re.IGNORECASE | |
| ) | |
| _keyword_heading_re = re.compile( | |
| r"^(Abstract|References?|Acknowledgments?|Related\s+Work" | |
| r"|Conclusion|Introduction|Appendix|Discussion" | |
| r"|Experiments?|Evaluation|Methodology|Datasets?)" | |
| r"(\b.{0,40})?$", # allow up to 40 extra chars (subtitle words) | |
| re.IGNORECASE, | |
| ) | |
| # ββ Pass 2: build Markdown βββββββββββββββββββββββββββββββββββββββββββββββ | |
| lines_out: list[str] = [] | |
| for page_num, page in enumerate(doc): | |
| page_height = page.rect.height | |
| blocks = page.get_text("dict")["blocks"] | |
| for block in blocks: | |
| if block.get("type") != 0: | |
| continue | |
| block_lines: list[str] = [] | |
| for line in block.get("lines", []): | |
| line_text_parts: list[str] = [] | |
| dominant_size = 0.0 | |
| for span in line.get("spans", []): | |
| txt = span["text"] | |
| if not txt.strip(): | |
| line_text_parts.append(txt) | |
| continue | |
| sz = span["size"] | |
| if sz > dominant_size: | |
| dominant_size = sz | |
| line_text_parts.append(txt) | |
| line_text = "".join(line_text_parts).strip() | |
| if not line_text: | |
| continue | |
| # Skip very short noise (page numbers, line numbers < 4 chars) | |
| if len(line_text) <= 3 and line_text.isdigit(): | |
| continue | |
| # Skip header/footer areas (top/bottom 5% of page) | |
| block_y = block["bbox"][1] | |
| if block_y < page_height * 0.05 or block_y > page_height * 0.95: | |
| continue | |
| # Determine heading level | |
| if dominant_size >= h1_threshold and page_num == 0: | |
| # Title: large font on first page only, must be > 8 chars | |
| if len(line_text) > 8: | |
| line_text = f"# {line_text}" | |
| elif dominant_size >= h2_threshold and ( | |
| _numbered_heading_re.match(line_text) | |
| or _keyword_heading_re.match(line_text) | |
| ): | |
| # Section heading: larger font AND looks like an ACL section title. | |
| # Combining both conditions avoids marking wrapped body lines as headings. | |
| line_text = f"## {line_text}" | |
| block_lines.append(line_text) | |
| if block_lines: | |
| lines_out.append("\n".join(block_lines)) | |
| doc.close() | |
| return "\n\n".join(lines_out) | |
| # ββ Markdown β PeerRead metadata conversion ββββββββββββββββββββββββββββββββββββ | |
| def _split_sections(md: str) -> list[dict]: | |
| """Split Markdown into PeerRead-style [{heading, text}] blocks. | |
| Matches the ground-truth convention: every heading yields a section even | |
| when its body is empty (PeerRead keeps heading-only sections, e.g. a parent | |
| "3" before "3.1"). A leading headingless block is only kept if it has text. | |
| """ | |
| lines = md.splitlines() | |
| sections: list[dict] = [] | |
| current_heading: str | None = None | |
| current_lines: list[str] = [] | |
| def _flush() -> None: | |
| text = "\n".join(current_lines).strip() | |
| # Headed sections are always emitted (text may be ""); a headingless | |
| # preamble is emitted only when it actually carries text. | |
| if current_heading is not None or text: | |
| sections.append({"heading": current_heading, "text": text}) | |
| for line in lines: | |
| m = re.match(r"^(#{1,4})\s+(.+)", line) | |
| if m: | |
| _flush() | |
| current_heading = m.group(2).strip() | |
| current_lines = [] | |
| else: | |
| current_lines.append(line) | |
| _flush() | |
| return sections | |
| def _extract_title(md: str) -> str: | |
| # H1 first (PyMuPDF provider) | |
| for line in md.splitlines(): | |
| m = re.match(r"^#\s+(.+)", line) | |
| if m: | |
| return m.group(1).strip() | |
| # Fallback: first H2 that doesn't look like a section heading (Docling uses H2 for titles) | |
| _section_re = re.compile(r"^(Abstract|Introduction|\d+[\s.])", re.IGNORECASE) | |
| for line in md.splitlines(): | |
| m = re.match(r"^##\s+(.+)", line) | |
| if m: | |
| heading = m.group(1).strip() | |
| if not _section_re.match(heading): | |
| return heading | |
| return "" | |
| def _extract_abstract(sections: list[dict]) -> str: | |
| for s in sections: | |
| if s["heading"] and re.search(r"\babstract\b", s["heading"], re.IGNORECASE): | |
| text = s["text"] | |
| # If the abstract section absorbed body text (> 2 000 chars) it means | |
| # subsequent section headings weren't detected. Return only the first | |
| # paragraph, which is the actual abstract. | |
| if len(text) > 2000: | |
| first_para = re.split(r"\n\n+", text)[0].strip() | |
| return first_para if len(first_para) > 50 else text[:2000] | |
| return text | |
| # Fallback: first headingless block long enough to be an abstract | |
| for s in sections: | |
| if not s["heading"] and len(s["text"]) > 150: | |
| return s["text"] | |
| return "" | |
| def _extract_authors(md: str) -> list[str]: | |
| """ | |
| Authors typically appear between the H1 title and the first section heading, | |
| as short lines of capitalized names (not affiliations or emails). | |
| """ | |
| lines = md.splitlines() | |
| in_preamble = False | |
| preamble: list[str] = [] | |
| for line in lines: | |
| if re.match(r"^#\s+", line): | |
| in_preamble = True | |
| continue | |
| if in_preamble: | |
| if re.match(r"^#{1,4}\s+", line): | |
| break | |
| preamble.append(line.strip()) | |
| authors: list[str] = [] | |
| for line in preamble: | |
| if not line: | |
| continue | |
| # Skip emails, URLs, affiliations | |
| if "@" in line or re.search(r"http|university|institute|department|lab\b|school", line, re.IGNORECASE): | |
| continue | |
| if len(line) > 100: | |
| continue | |
| # Must look like a name: starts capital, only letters/spaces/commas/hyphens/dots | |
| if re.match(r"^[A-Z][a-zA-Z\s,.\-]+$", line): | |
| parts = re.split(r",\s*|\s+and\s+", line) | |
| authors.extend(p.strip() for p in parts if p.strip()) | |
| return authors | |
| def _extract_year(md: str) -> int | None: | |
| # Only look in the preamble (before first section heading) to avoid | |
| # picking up years from citations in body text. | |
| preamble_end = len(md) | |
| for m in re.finditer(r"^#{1,4}\s+", md, re.MULTILINE): | |
| if m.start() > 0: | |
| preamble_end = m.start() | |
| break | |
| preamble = md[:min(preamble_end, 1500)] | |
| for m in re.finditer(r"\b(20[0-2]\d|19\d\d)\b", preamble): | |
| year = int(m.group(1)) | |
| if 1990 <= year <= 2030: | |
| return year | |
| return None | |
| # ββ Reference parsing ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ | |
| def _last_name(author: str) -> str: | |
| """Last whitespace token of an author name, KEEPING any trailing period. | |
| GROBID's citeRegEx uses the raw token, so "Marianna Apidianaki." β "Apidianaki." | |
| and "Marco Baroni" β "Baroni" (the period only survives on the final author of | |
| the list, which is exactly how the ground truth builds it). | |
| """ | |
| toks = author.split() | |
| return toks[-1] if toks else "" | |
| def _build_cite_regex(authors: list[str], year: int | None) -> tuple[str, str]: | |
| """Return (citeRegEx, shortCiteRegEx) matching the GROBID author-year style.""" | |
| if not authors or year is None: | |
| return "", "" | |
| def _name_pat(name: str) -> str: | |
| # Parsed last names may keep a trailing period from bibliography text. | |
| # Body citations usually do not, so make that period optional and escape | |
| # the rest of the name for regex safety. | |
| return re.escape(name.rstrip(".")).replace(r"\ ", r"\s+") + r"\.?" | |
| ln1 = _last_name(authors[0]) | |
| if not ln1: | |
| return "", "" | |
| if len(authors) >= 3: | |
| short = rf"{_name_pat(ln1)}\s+et\s+al\.?" | |
| elif len(authors) == 2: | |
| # The second author is often mangled by PDF text extraction | |
| # (e.g. "M` arquez", "PadΒ΄ o"). For mention detection, first-author | |
| # + "and <short name phrase>" + year is robust enough; the metric key | |
| # is first-author/year anyway. | |
| short = rf"{_name_pat(ln1)}\s+and\s+[A-Z][^,();]{{1,60}}?" | |
| else: | |
| short = _name_pat(ln1) | |
| # Match both parenthetical citations ("Name, 2016") and narrative citations | |
| # ("Name (2016)" or "Name (2015, 2016)"). The latter is common in PeerRead | |
| # bodies and was previously under-counted as a missing mention. | |
| year_pat = rf"{year}[a-z]?" | |
| parenthetical = rf"{short},?\s+{year_pat}" | |
| narrative = rf"{short}\s*\([^)]*\b{year_pat}\b[^)]*\)" | |
| return rf"(?:{parenthetical}|{narrative})", short | |
| def _parse_reference(text: str) -> dict | None: | |
| """Parse one author-year reference entry β PeerRead reference dict. | |
| Expected ACL layout: "Authors. Year. Title. Venue." (also tolerates a | |
| leading "[N]"/"N."/bullet and a quoted title). | |
| """ | |
| text = re.sub(r"^[\[(]?\d+[\]).]?\s*", "", text).strip() # leading [N] / N. | |
| text = re.sub(r"^[-*β’Β·]\s*", "", text).strip() # leading bullet | |
| if len(text) < 10: | |
| return None | |
| # Author/year boundary = first 19xx/20xx, preferring one followed by a period. | |
| year_m = ( | |
| re.search(r"\b((?:19|20)\d\d)[a-z]?\.", text) | |
| or re.search(r"\b((?:19|20)\d\d)[a-z]?\b", text) | |
| ) | |
| if not year_m: | |
| return None | |
| year = int(year_m.group(1)) | |
| authors_str = text[: year_m.start()].rstrip() | |
| rest = text[year_m.end():].strip() | |
| authors: list[str] = [] | |
| if authors_str: | |
| parts = re.split(r",\s*(?:and\s+)?|\s+and\s+", authors_str) | |
| authors = [p.strip() for p in parts if p.strip() and len(p.strip()) < 60] | |
| # Title / venue from the remainder ("Title. Venue"). | |
| title = venue = "" | |
| qm = re.search(r'["β](.+?)["β]', rest) | |
| if qm: | |
| title = qm.group(1).strip() | |
| venue = rest[qm.end():].strip(' .,') | |
| elif rest: | |
| tparts = rest.split(". ", 1) | |
| title = tparts[0].strip().rstrip(".") | |
| venue = tparts[1].strip() if len(tparts) > 1 else "" | |
| cite_regex, short_cite = _build_cite_regex(authors, year) | |
| return { | |
| "title": title, | |
| "author": authors, | |
| "venue": venue, | |
| "citeRegEx": cite_regex, | |
| "shortCiteRegEx": short_cite, | |
| "year": year, | |
| } | |
| def _split_ref_entries(block: str) -> list[str]: | |
| """Split a References block into per-entry strings. | |
| Docling usually emits each bibliography item as its own block (blank-line | |
| separated). If that doesn't hold, fall back to grouping wrapped lines: a new | |
| entry begins at an author-like line once the buffer already carries a year. | |
| """ | |
| block = block.strip() | |
| year_re = re.compile(r"\b(?:19|20)\d\d") | |
| paras = [p.strip() for p in re.split(r"\n\s*\n", block) if p.strip()] | |
| if len(paras) >= 2 and sum(1 for p in paras if year_re.search(p)) >= len(paras) * 0.6: | |
| return paras | |
| lines = [ln.strip() for ln in block.splitlines() if ln.strip()] | |
| start_re = re.compile(r"^[A-Z][A-Za-z.'\-]+(?:,|\s+[A-Z])") # "Surname," / "Surname Initial" | |
| entries: list[str] = [] | |
| buf: list[str] = [] | |
| for ln in lines: | |
| if buf and year_re.search(" ".join(buf)) and start_re.match(ln): | |
| entries.append(" ".join(buf)) | |
| buf = [ln] | |
| else: | |
| buf.append(ln) | |
| if buf: | |
| entries.append(" ".join(buf)) | |
| return entries | |
| def _parse_references(md: str) -> list[dict]: | |
| """Extract and parse the References / Bibliography section.""" | |
| m = re.search( | |
| r"^#{1,4}\s+(?:References?|Bibliography)\s*$", md, re.MULTILINE | re.IGNORECASE | |
| ) | |
| if not m: | |
| return [] | |
| ref_block = md[m.end():] | |
| next_h = re.search(r"^#{1,4}\s+", ref_block, re.MULTILINE) # cut at next heading | |
| if next_h: | |
| ref_block = ref_block[: next_h.start()] | |
| refs = [] | |
| for entry in _split_ref_entries(ref_block): | |
| r = _parse_reference(entry) | |
| if r: | |
| refs.append(r) | |
| return refs | |
| def _sentence_context(text: str, start: int, end: int) -> tuple[str, int, int]: | |
| """Sentence containing [start, end), plus the match offsets within it. | |
| GROBID's referenceMentions use sentence-bounded context (not a fixed window). | |
| """ | |
| prev = list(re.finditer(r"[.!?]\s+", text[:start])) | |
| ctx_start = prev[-1].end() if prev else 0 | |
| nxt = re.search(r"[.!?]\s+", text[end:]) | |
| ctx_end = end + nxt.start() + 1 if nxt else len(text) | |
| ctx = text[ctx_start:ctx_end] | |
| lead = len(ctx) - len(ctx.lstrip()) | |
| ctx = ctx.strip() | |
| return ctx, start - ctx_start - lead, end - ctx_start - lead | |
| def _extract_reference_mentions(sections: list[dict], references: list[dict]) -> list[dict]: | |
| """Scan body text for author-year citations and build referenceMentions. | |
| Each reference's citeRegEx is matched against every body section (the | |
| References section itself is skipped). Mentions are emitted in document order | |
| with sentence-bounded context + offsets, matching the PeerRead schema. | |
| """ | |
| compiled: list[tuple[int, tuple[str, int] | None, "re.Pattern[str]"]] = [] | |
| for i, r in enumerate(references): | |
| pat = r.get("citeRegEx") or "" | |
| if not pat: | |
| continue | |
| authors = r.get("author") or [] | |
| toks = authors[0].split() if authors else [] | |
| ln = toks[-1].strip(".,").lower() if toks else "" | |
| year = r.get("year") | |
| key = (ln, year) if ln and isinstance(year, int) else None | |
| try: | |
| compiled.append((i, key, re.compile(pat))) | |
| except re.error: | |
| continue | |
| mentions: list[dict] = [] | |
| for section in sections: | |
| heading = (section.get("heading") or "").strip() | |
| if re.match(r"references?$", heading, re.IGNORECASE): | |
| continue | |
| text = section.get("text", "") | |
| if not text: | |
| continue | |
| found: list[tuple[int, int, int, tuple[str, int] | None]] = [] | |
| seen_spans: set[tuple[int, int, tuple[str, int] | None]] = set() | |
| for ref_id, key, rx in compiled: | |
| for mm in rx.finditer(text): | |
| span_key = (mm.start(), mm.end(), key) | |
| if span_key in seen_spans: | |
| continue | |
| seen_spans.add(span_key) | |
| found.append((mm.start(), mm.end(), ref_id, key)) | |
| found.sort() # document order | |
| for s_, e_, ref_id, _key in found: | |
| ctx, off_s, off_e = _sentence_context(text, s_, e_) | |
| mentions.append({ | |
| "referenceID": ref_id, | |
| "context": ctx, | |
| "startOffset": off_s, | |
| "endOffset": off_e, | |
| }) | |
| return mentions | |
| def _refresh_reference_regexes(references: list[dict]) -> list[dict]: | |
| """Rebuild citation regexes with the current parser logic.""" | |
| refreshed: list[dict] = [] | |
| for ref in references: | |
| r = dict(ref) | |
| cite_regex, short_cite = _build_cite_regex(r.get("author") or [], r.get("year")) | |
| if cite_regex: | |
| r["citeRegEx"] = cite_regex | |
| r["shortCiteRegEx"] = short_cite | |
| refreshed.append(r) | |
| return refreshed | |
| def _sections_for_mention_eval(meta: dict) -> list[dict]: | |
| """Sections used when rebuilding mentions from saved PeerRead-style JSON. | |
| `md_to_peerread` extracts mentions before dropping Abstract from `sections`, | |
| but saved JSON keeps the abstract separately as `abstractText`. Rebuild from | |
| both sides with the same source text so the metric is symmetric. | |
| """ | |
| sections: list[dict] = [] | |
| abstract = (meta.get("abstractText") or "").strip() | |
| if abstract: | |
| sections.append({"heading": "Abstract", "text": abstract}) | |
| sections.extend(meta.get("sections", [])) | |
| return sections | |
| def _normalize_mentions_for_eval(meta: dict) -> dict: | |
| """Rebuild references' regexes and mentions with the current parser logic.""" | |
| out = dict(meta) | |
| out["references"] = _refresh_reference_regexes(out.get("references", [])) | |
| out["referenceMentions"] = _extract_reference_mentions( | |
| _sections_for_mention_eval(out), out["references"] | |
| ) | |
| return out | |
| # ββ Top-level converter ββββββββββββββββββββββββββββββββββββββββββββββββββββββββ | |
| def md_to_peerread(md: str, paper_id: str, provider: str = "docling") -> dict: | |
| """Convert Markdown (from any provider) to PeerRead metadata JSON.""" | |
| sections = _split_sections(md) | |
| title = _extract_title(md) | |
| abstract = _extract_abstract(sections) | |
| authors = _extract_authors(md) | |
| year = _extract_year(md) | |
| references = _parse_references(md) | |
| mentions = _extract_reference_mentions(sections, references) | |
| # Match the ground-truth `sections` convention: it excludes the title block, | |
| # the Abstract (lives only in abstractText) and the References bibliography | |
| # (lives only in references[]). Appendices after the bibliography are kept. | |
| def _is_excluded(heading: str | None) -> bool: | |
| if heading is None: | |
| return False | |
| h = heading.strip() | |
| if h == title: | |
| return True | |
| if re.search(r"\babstract\b", h, re.IGNORECASE): | |
| return True | |
| if re.match(r"references?$", h, re.IGNORECASE): | |
| return True | |
| return False | |
| body_sections = [s for s in sections if not _is_excluded(s.get("heading"))] | |
| return { | |
| "name": f"{paper_id}.pdf", | |
| "metadata": { | |
| "source": provider, | |
| "title": title, | |
| "authors": authors, | |
| "emails": [], | |
| "sections": body_sections, | |
| "references": references, | |
| "referenceMentions": mentions, | |
| "year": year, | |
| "abstractText": abstract, | |
| "creator": provider.capitalize(), | |
| }, | |
| } | |
| # ββ OCR runner βββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ | |
| async def _get_markdown(pdf_bytes: bytes, provider: str) -> str: | |
| if provider == "local": | |
| return _pdf2md_local(pdf_bytes) | |
| if provider == "docling-granite": | |
| return await _pdf2md_docling_server(pdf_bytes) | |
| if provider == "docling-granite-local": | |
| return _pdf2md_docling_granite_local(pdf_bytes) | |
| if provider == "docling-standard-local": | |
| return _pdf2md_docling_standard_local(pdf_bytes) | |
| raise ValueError(f"Unknown provider: {provider!r}") | |
| async def process_pdf( | |
| pdf_path: Path, out_dir: Path, provider: str, overwrite: bool = False | |
| ) -> None: | |
| paper_id = pdf_path.stem | |
| out_path = out_dir / f"{paper_id}.pdf.json" | |
| if out_path.exists() and not overwrite: | |
| print(f"[skip] {paper_id} (already done β use --overwrite to redo)") | |
| return | |
| label = { | |
| "local": "locally (PyMuPDF)", | |
| "docling-granite": "Docling Granite VLM pipeline (Kaggle server)", | |
| "docling-granite-local": "Docling Granite VLM pipeline (local)", | |
| "docling-standard-local": "Docling standard pipeline + formula enrichment (local)", | |
| }.get(provider, provider) | |
| print(f"[ .. ] {paper_id} processing via {label} ...", end="", flush=True) | |
| try: | |
| md = await _get_markdown(pdf_path.read_bytes(), provider) | |
| result = md_to_peerread(md, paper_id, provider) | |
| out_path.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8") | |
| n_sec = len(result["metadata"]["sections"]) | |
| n_ref = len(result["metadata"]["references"]) | |
| print(f"\r[ ok ] {paper_id} {n_sec} sections {n_ref} refs -> {out_path.name}") | |
| except Exception as exc: | |
| import traceback | |
| print(f"\r[ERR] {paper_id} {exc}") | |
| traceback.print_exc() | |
| async def run_ocr(pdf_dir: Path, out_dir: Path, provider: str, overwrite: bool) -> None: | |
| # Resolve to absolute path before any library call that might change CWD | |
| out_dir = out_dir.resolve() | |
| pdf_dir = pdf_dir.resolve() | |
| out_dir.mkdir(parents=True, exist_ok=True) | |
| pdfs = sorted(pdf_dir.glob("*.pdf")) | |
| if not pdfs: | |
| sys.exit(f"No PDFs found in {pdf_dir}") | |
| print(f"Processing {len(pdfs)} PDFs from {pdf_dir} [provider={provider}]\n") | |
| for pdf in pdfs: | |
| await process_pdf(pdf, out_dir, provider, overwrite) | |
| print(f"\nDone. Output in {out_dir}/") | |
| # ββ Comparison report ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ | |
| def _normalize_for_eval(text: str) -> str: | |
| """Canonicalize text before similarity scoring, applied to BOTH GT and pred. | |
| The ground truth (GROBID) carries noise that Docling cleanly omits β margin | |
| line numbers, page numbers, mojibake (ΟβΓ, ββΓ’1) and odd formula spacing. | |
| Comparing raw text unfairly penalizes the cleaner output, so we fold both | |
| sides to a common form: drop standalone numeric lines, fold non-ASCII | |
| (mojibake and real unicode symbols alike), lowercase, collapse whitespace. | |
| """ | |
| lines = [ | |
| ln for ln in text.splitlines() | |
| if not re.fullmatch(r"\s*\d{1,4}[.)]?\s*", ln) # pure margin/page numbers | |
| ] | |
| s = " ".join(lines) | |
| s = unicodedata.normalize("NFKD", s).encode("ascii", "ignore").decode("ascii") | |
| s = re.sub(r"\s+", " ", s).strip().lower() | |
| return s | |
| def _body_text(meta: dict) -> str: | |
| return "\n".join(s.get("text", "") for s in meta.get("sections", [])) | |
| def _ref_keys(meta: dict) -> set[tuple[str, int]]: | |
| """(first-author last name, year) keys β robust to formatting differences.""" | |
| keys: set[tuple[str, int]] = set() | |
| for r in meta.get("references", []): | |
| authors = r.get("author") or [] | |
| toks = authors[0].split() if authors else [] | |
| ln = toks[-1].strip(".,").lower() if toks else "" | |
| year = r.get("year") | |
| if ln and isinstance(year, int): | |
| keys.add((ln, year)) | |
| return keys | |
| def _loose(text: str) -> str: | |
| """Space-insensitive canonical form (applied to BOTH sides): drop margin/page | |
| numbers, fold non-ASCII (GT mojibake ΟβΓ), lowercase, and REMOVE ALL spaces. | |
| OCR often inserts spaces between characters (e.g. "h i -1", "s ( c ) k"); the | |
| ground truth spaces them differently. Removing whitespace makes scoring lenient | |
| about spacing so it reflects actual character content, not layout. | |
| """ | |
| lines = [ | |
| ln for ln in text.splitlines() | |
| if not re.fullmatch(r"\s*\d{1,4}[.)]?\s*", ln) | |
| ] | |
| s = " ".join(lines) | |
| s = unicodedata.normalize("NFKD", s).encode("ascii", "ignore").decode("ascii").lower() | |
| return re.sub(r"\s+", "", s) | |
| def _sim_loose(a: str, b: str) -> float: | |
| """Character similarity (β, 0β1) ignoring whitespace differences.""" | |
| la, lb = _loose(a), _loose(b) | |
| if not la and not lb: | |
| return 1.0 | |
| return SequenceMatcher(None, la, lb, autojunk=False).ratio() | |
| def _tokens_for_eval(text: str) -> list[str]: | |
| """Word/number tokens after the same loose normalization used for scoring.""" | |
| return re.findall(r"[a-z0-9]+(?:[.-][a-z0-9]+)*", _normalize_for_eval(text)) | |
| def _multiset_prf(gt_items: list, pred_items: list) -> tuple[float, float, float]: | |
| """Precision/recall/F1 over multisets (both empty β perfect).""" | |
| from collections import Counter | |
| g, p = Counter(gt_items), Counter(pred_items) | |
| if not g and not p: | |
| return 1.0, 1.0, 1.0 | |
| inter = sum((g & p).values()) | |
| rec = inter / sum(g.values()) if g else 0.0 | |
| prec = inter / sum(p.values()) if p else 0.0 | |
| f1 = 2 * prec * rec / (prec + rec) if (prec + rec) else 0.0 | |
| return prec, rec, f1 | |
| def _classify_lines(text: str) -> tuple[list[str], list[str]]: | |
| """Split body text into (prose lines, table-ish lines). | |
| A table-ish line is number-dense (β₯3 numeric tokens and β₯40% of tokens | |
| numeric) β GROBID flattens tables into exactly this shape. | |
| """ | |
| prose, table = [], [] | |
| for ln in text.splitlines(): | |
| s = ln.strip() | |
| if not s or re.fullmatch(r"\d{1,4}[.)]?", s): | |
| continue | |
| toks = s.split() | |
| nums = [t for t in toks if re.search(r"\d", t)] | |
| if toks and len(nums) >= 3 and len(nums) >= 0.4 * len(toks): | |
| table.append(s) | |
| else: | |
| prose.append(s) | |
| return prose, table | |
| _CLEAN_CHAR_RE = re.compile(r"[A-Za-z0-9 .,%()\[\];:/+\-#|~]") | |
| def _is_junk_line(line: str) -> bool: | |
| """True for a GROBID line that is really a mis-OCR'd figure/CJK block. | |
| PeerRead's ground truth sometimes flattens a figure, CJK caption or score | |
| glyph into the body as mojibake dominated by punctuation ($ & ! " ζΆ β¦). | |
| These leak into the table/formula buckets and unfairly drag scores down, so | |
| numeric/formula scoring skips them. Threshold: a genuine data row is β₯80% | |
| "clean" chars (letters, digits, %, brackets, common punctuation); mojibake | |
| rows fall below it. | |
| """ | |
| s = line.strip() | |
| if not s: | |
| return True | |
| clean = sum(1 for c in s if _CLEAN_CHAR_RE.match(c)) | |
| return clean / len(s) < 0.80 | |
| def _data_numbers(text: str) -> list[str]: | |
| """Meaningful numeric values: decimals, percentages, or β₯2-digit/comma counts. | |
| Drops lone single digits β GROBID renders checkmarks/daggers as "7" and | |
| section refs as "3", which are noise rather than reported table data. | |
| """ | |
| out: list[str] = [] | |
| for t in re.findall(r"\d[\d.,%]*", text): | |
| core = t.rstrip(".,%") | |
| if "." in core or "%" in t or "," in core or len(core) >= 2: | |
| out.append(core) | |
| return out | |
| def _heading_seq(meta: dict) -> list[str]: | |
| headings = [] | |
| for s in meta.get("sections", []): | |
| if not s.get("heading"): | |
| continue | |
| h = _normalize_for_eval(s["heading"]) | |
| # Output can keep submission boilerplate that GT dropped. That should not | |
| # count against reading-order preservation. | |
| if h and h != "anonymous acl submission": | |
| headings.append(h) | |
| return headings | |
| def _lcs_len(a: list, b: list) -> int: | |
| if not a or not b: | |
| return 0 | |
| prev = [0] * (len(b) + 1) | |
| for x in a: | |
| cur = [0] * (len(b) + 1) | |
| for j, y in enumerate(b): | |
| cur[j + 1] = prev[j] + 1 if x == y else max(prev[j + 1], cur[j]) | |
| prev = cur | |
| return prev[-1] | |
| def _mention_keys(meta: dict) -> list[tuple[str, int]]: | |
| """Resolve each referenceMention to its (last-name, year) via referenceID.""" | |
| refs = meta.get("references", []) | |
| keys: list[tuple[str, int]] = [] | |
| for mn in meta.get("referenceMentions", []): | |
| rid = mn.get("referenceID") | |
| if isinstance(rid, int) and 0 <= rid < len(refs): | |
| authors = refs[rid].get("author") or [] | |
| toks = authors[0].split() if authors else [] | |
| ln = toks[-1].strip(".,").lower() if toks else "" | |
| year = refs[rid].get("year") | |
| if ln and isinstance(year, int): | |
| keys.append((ln, year)) | |
| return keys | |
| def _component_scores(gt: dict, pred: dict) -> dict: | |
| """Per-component metrics for one paper (see legend in run_compare).""" | |
| gt_body, pred_body = _body_text(gt), _body_text(pred) | |
| gt_prose, gt_table = _classify_lines(gt_body) | |
| pr_prose, pr_table = _classify_lines(pred_body) | |
| gt_prose_text = "\n".join(gt_prose) | |
| pr_prose_text = "\n".join(pr_prose) | |
| text_sim = _sim_loose(gt_prose_text, pr_prose_text) # Text char diagnostic β | |
| text_p, text_r, text_f = _multiset_prf( | |
| _tokens_for_eval(gt_prose_text), _tokens_for_eval(pr_prose_text) | |
| ) | |
| # Table: coverage of the GROBID table's DATA VALUES in the prediction. | |
| # GROBID and Docling flatten tables differently (space- vs dot-joined), so | |
| # classifying pred lines independently is unreliable β it scored a paper | |
| # 0.000 although every number was captured. Instead take the data values | |
| # GROBID placed in table-ish rows (minus mojibake figure rows) and measure | |
| # how many appear ANYWHERE in the prediction body: distinct-value recall, | |
| # robust to flatten/spacing and to GROBID's spurious duplicated rows. | |
| gt_table_clean = [l for l in gt_table if not _is_junk_line(l)] | |
| pr_table_clean = [l for l in pr_table if not _is_junk_line(l)] | |
| table_sim = _sim_loose("\n".join(gt_table_clean), "\n".join(pr_table_clean)) # diagnostic β | |
| gt_dnums = set(_data_numbers("\n".join(gt_table_clean))) | |
| pred_dnums_all = set(_data_numbers(pred_body)) | |
| tbl_num_cov = len(gt_dnums & pred_dnums_all) / len(gt_dnums) if gt_dnums else 1.0 | |
| # multiset P/R/F on the table buckets, kept as diagnostics only | |
| tbl_num_p, tbl_num_r, tbl_num_f = _multiset_prf( | |
| _data_numbers("\n".join(gt_table_clean)), _data_numbers("\n".join(pr_table_clean)) | |
| ) | |
| # Formula: equation recall via "=" count, on non-junk lines only so a mojibake | |
| # "=" inside a mis-OCR'd figure row can't inflate the GT denominator. | |
| eg = sum(ln.count("=") for ln in gt_body.splitlines() if not _is_junk_line(ln)) | |
| ep = sum(ln.count("=") for ln in pred_body.splitlines() if not _is_junk_line(ln)) | |
| formula_rec = min(ep, eg) / eg if eg else 1.0 | |
| gh, ph = _heading_seq(gt), _heading_seq(pred) # Order LCS β | |
| order = _lcs_len(gh, ph) / len(gh) if gh else (1.0 if not ph else 0.0) | |
| ref_p, ref_r, ref_f = _multiset_prf(list(_ref_keys(gt)), list(_ref_keys(pred))) | |
| men_p, men_r, men_f = _multiset_prf(_mention_keys(gt), _mention_keys(pred)) | |
| return { | |
| "text_sim": text_sim, "text_p": text_p, "text_r": text_r, "text_f": text_f, | |
| "table_sim": table_sim, | |
| "tbl_num_cov": tbl_num_cov, | |
| "tbl_num_p": tbl_num_p, "tbl_num_r": tbl_num_r, "tbl_num_f": tbl_num_f, | |
| "formula": formula_rec, "order": order, | |
| "ref_p": ref_p, "ref_r": ref_r, "ref_f": ref_f, | |
| "men_p": men_p, "men_r": men_r, "men_f": men_f, | |
| } | |
| def run_compare(gt_dir: Path, pred_dir: Path) -> None: | |
| gt_files = {p.name: p for p in gt_dir.glob("*.pdf.json")} | |
| pred_files = {p.name: p for p in pred_dir.glob("*.pdf.json")} | |
| common = sorted(gt_files.keys() & pred_files.keys()) | |
| if not common: | |
| sys.exit(f"No matching files between {gt_dir} and {pred_dir}") | |
| rows: list[tuple[str, dict]] = [] | |
| for name in common: | |
| gt = json.loads(gt_files[name].read_text(encoding="utf-8"))["metadata"] | |
| pred = json.loads(pred_files[name].read_text(encoding="utf-8"))["metadata"] | |
| # Apply parser improvements symmetrically without forcing a full Docling | |
| # re-run. This avoids comparing normalized pred mentions against older | |
| # GROBID mention extraction quirks from the saved GT JSON. | |
| gt = _normalize_mentions_for_eval(gt) | |
| pred = _normalize_mentions_for_eval(pred) | |
| rows.append((name.split(".")[0], _component_scores(gt, pred))) | |
| n = len(rows) | |
| avg = lambda k: sum(s[k] for _, s in rows) / n | |
| print(f"Comparing {n} paper(s) (all scores 0β1, HIGHER = better, space-insensitive)\n") | |
| # ββ Headline: one easy score per data type βββββββββββββββββββββββββββββββ | |
| sh = (f"{'ID':<6} {'Text':>6} {'Table':>6} {'Formula':>8} " | |
| f"{'Order':>6} {'Refs':>6} {'Mentions':>9}") | |
| print(sh) | |
| print("-" * len(sh)) | |
| for pid, s in rows: | |
| print(f"{pid:<6} {s['text_r']:>6.3f} {s['tbl_num_cov']:>6.3f} {s['formula']:>8.3f} " | |
| f"{s['order']:>6.3f} {s['ref_r']:>6.3f} {s['men_r']:>9.3f}") | |
| print("-" * len(sh)) | |
| print(f"{'AVG':<6} {avg('text_r'):>6.3f} {avg('tbl_num_cov'):>6.3f} {avg('formula'):>8.3f} " | |
| f"{avg('order'):>6.3f} {avg('ref_r'):>6.3f} {avg('men_r'):>9.3f}") | |
| # ββ Precision / Recall detail ββββββββββββββββββββββββββββββββββββββββββββ | |
| dh = (f"{'ID':<6} {'TxtP':>6} {'TxtF':>6} {'TxtChr':>6} " | |
| f"{'TblP':>6} {'TblR':>6} {'TblF':>6} {'TblTxt':>6} " | |
| f"{'RefP':>6} {'RefR':>6} {'MenP':>6} {'MenR':>6}") | |
| print("\n=== Detail (text/table diagnostics, reference & mention P/R) ===") | |
| print(dh) | |
| print("-" * len(dh)) | |
| for pid, s in rows: | |
| print(f"{pid:<6} {s['text_p']:>6.3f} {s['text_f']:>6.3f} {s['text_sim']:>6.3f} " | |
| f"{s['tbl_num_p']:>6.3f} {s['tbl_num_r']:>6.3f} " | |
| f"{s['tbl_num_f']:>6.3f} {s['table_sim']:>6.3f} " | |
| f"{s['ref_p']:>6.3f} {s['ref_r']:>6.3f} {s['men_p']:>6.3f} {s['men_r']:>6.3f}") | |
| print("-" * len(dh)) | |
| print(f"{'AVG':<6} {avg('text_p'):>6.3f} {avg('text_f'):>6.3f} {avg('text_sim'):>6.3f} " | |
| f"{avg('tbl_num_p'):>6.3f} {avg('tbl_num_r'):>6.3f} " | |
| f"{avg('tbl_num_f'):>6.3f} {avg('table_sim'):>6.3f} " | |
| f"{avg('ref_p'):>6.3f} {avg('ref_r'):>6.3f} {avg('men_p'):>6.3f} {avg('men_r'):>6.3f}") | |
| print("\nText β : GT-token recall on prose (does not punish useful extra extraction).") | |
| print("Table β : distinct GT table data-values found in pred (junk rows filtered).") | |
| print("Formula β : formula recall (decoded equations captured vs GT).") | |
| print("Order β : GT heading-order recall (LCS / #GT headings).") | |
| print("Refs β : reference recall by (last-name, year).") | |
| print("Mentions β : citationβreference linking recall.") | |
| print("TxtChr/TblTxt β : old char-similarity diagnostics only.") | |
| # ββ Entry point ββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββββ | |
| def main() -> None: | |
| parser = argparse.ArgumentParser( | |
| description="Test OCR against PeerRead ACL-2017 test set" | |
| ) | |
| parser.add_argument("--pdf-dir", default="test/pdfs", help="Input PDF directory") | |
| parser.add_argument("--out-dir", default="test/docling_output", help="Output JSON directory") | |
| parser.add_argument( | |
| "--provider", | |
| default="docling-granite", | |
| choices=[ | |
| "docling-granite", # Granite VLM via Kaggle server | |
| "docling-granite-local", # Granite VLM locally (needs GPU; slow on CPU) | |
| "docling-standard-local", # standard pipeline + formula enrichment (decodes formulas) | |
| "local", # PyMuPDF baseline | |
| ], | |
| help="docling-standard-local=decodes formulas (recommended), docling-granite=VLM via Kaggle, local=PyMuPDF baseline", | |
| ) | |
| parser.add_argument("--overwrite", action="store_true", help="Re-process already-done files") | |
| parser.add_argument( | |
| "--compare", | |
| action="store_true", | |
| help="Compare test/docling_output/ vs test/parsed_pdfs/ (ground truth)", | |
| ) | |
| parser.add_argument( | |
| "--gt-dir", | |
| default="test/parsed_pdfs", | |
| help="Ground truth directory (for --compare)", | |
| ) | |
| args = parser.parse_args() | |
| if args.compare: | |
| run_compare(Path(args.gt_dir), Path(args.out_dir)) | |
| else: | |
| asyncio.run(run_ocr(Path(args.pdf_dir), Path(args.out_dir), args.provider, args.overwrite)) | |
| if __name__ == "__main__": | |
| main() | |