Add data package to HF Space

src/paper2lab/data/pdf_loader.py

@@ -0,0 +1,974 @@
+"""
+pdf_loader.py — Paper2Lab universal section-aware PDF ingestion.
+
+Final extraction-layer guarantees:
+- Field-agnostic section detection across ML, NLP, CV, biomedical, physics,
+  education, social-science, economics, and interdisciplinary papers.
+- Multi-signal heading confidence scoring: lexical headings, numbered/roman
+  headings, ALL-CAPS headings, font size, bold text, and vertical spacing.
+- General anti-noise rules: boilerplate, metric-only headings, table-cell-like
+  headings, reference items, and fragmented tables.
+- Clean body text excludes References/Bibliography, Appendix/Supplementary, and
+  boilerplate so downstream extraction does not get polluted.
+- raw_text and all_sections are preserved for traceability/debugging.
+"""
+
+from __future__ import annotations
+
+import re
+from collections import Counter, defaultdict
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Set, Tuple
+
+import fitz  # PyMuPDF
+
+
+class PDFIngestionError(Exception):
+    pass
+
+
+# ---------------------------------------------------------------------------
+# Section taxonomy
+# ---------------------------------------------------------------------------
+
+SECTION_KEYWORDS: Set[str] = {
+    # Universal/front matter
+    "abstract", "keywords", "introduction", "background", "overview",
+    "motivation", "problem statement", "problem formulation",
+
+    # Related work / literature
+    "related work", "related works", "prior work", "prior art",
+    "literature review", "literature survey", "state of the art",
+
+    # Theory / framework
+    "theoretical background", "theory", "framework", "preliminaries",
+    "notation", "problem definition", "mathematical background",
+
+    # Methodology
+    "materials and methods", "material and methods", "methods", "method",
+    "methodology", "approach", "proposed approach", "proposed method",
+    "model", "model architecture", "architecture", "system design",
+    "system overview", "implementation", "implementation details",
+    "technical details", "training", "training details", "training procedure",
+    "experimental setup", "experimental settings", "experimental design",
+    "data collection", "data preprocessing", "data preparation",
+    "study design", "participants", "procedure",
+
+    # Experiments/results
+    "experiments", "experiment", "evaluation", "evaluations",
+    "results", "results and discussion", "results and analysis",
+    "empirical evaluation", "empirical results", "analysis",
+    "quantitative analysis", "qualitative analysis", "ablation",
+    "ablation study", "ablation studies", "case study", "case studies",
+
+    # Discussion/limitations/conclusion
+    "discussion", "limitations", "limitation", "future work",
+    "future directions", "conclusion", "conclusions", "summary",
+    "concluding remarks",
+
+    # Admin/back matter
+    "data availability", "code availability", "availability",
+    "ethics statement", "ethical considerations", "ethical approval",
+    "acknowledgements", "acknowledgments", "funding",
+    "conflict of interest", "competing interests", "author contributions",
+    "references", "bibliography", "works cited", "literature cited",
+    "appendix", "appendices", "supplementary material",
+    "supplementary information", "supplemental material", "supplementary",
+    "supplemental",
+}
+
+REFERENCE_SECTION_NAMES: Set[str] = {
+    "references", "bibliography", "works cited", "literature cited",
+}
+
+APPENDIX_SECTION_NAMES: Set[str] = {
+    "appendix", "appendices", "supplementary material",
+    "supplementary information", "supplemental material", "supplementary",
+    "supplemental",
+}
+
+BOILERPLATE_SECTION_NAMES: Set[str] = {
+    "acknowledgements", "acknowledgments", "author contributions",
+    "competing interests", "conflict of interest", "additional information",
+    "publisher's note", "open access", "correspondence",
+    "reprints and permissions", "funding", "ethics statement",
+    "ethical considerations", "ethical approval", "data availability",
+    "code availability", "availability",
+}
+
+ROLE_ALIASES: Dict[str, List[str]] = {
+    "front_matter": ["front matter"],
+    "abstract": ["abstract"],
+    "keywords": ["keywords"],
+    "introduction": [
+        "introduction", "overview", "motivation", "problem statement",
+        "problem formulation",
+    ],
+    "related_work": ["related work", "related works", "prior work", "prior art"],
+    "background": ["background", "preliminaries", "notation"],
+    "theory": [
+        "theoretical background", "theory", "framework",
+        "mathematical background", "problem definition",
+    ],
+    "methodology": [
+        "materials and methods", "material and methods", "methods", "method",
+        "methodology", "approach", "proposed approach", "proposed method",
+        "model", "model architecture", "architecture", "system design",
+        "system overview", "implementation", "implementation details",
+        "technical details", "training", "training details", "training procedure",
+        "experimental setup", "experimental settings", "experimental design",
+        "data collection", "data preprocessing", "data preparation",
+        "study design", "participants", "procedure",
+    ],
+    "experiments": [
+        "experiments", "experiment", "evaluation", "evaluations",
+        "empirical evaluation", "ablation", "ablation study", "ablation studies",
+        "case study", "case studies",
+    ],
+    "results": [
+        "results", "results and discussion", "results and analysis",
+        "empirical results", "quantitative analysis", "qualitative analysis",
+        "analysis",
+    ],
+    "discussion": ["discussion"],
+    "limitations": ["limitations", "limitation"],
+    "future_work": ["future work", "future directions"],
+    "conclusion": ["conclusion", "conclusions", "summary", "concluding remarks"],
+    "references": list(REFERENCE_SECTION_NAMES),
+    "appendix": list(APPENDIX_SECTION_NAMES),
+    "boilerplate": list(BOILERPLATE_SECTION_NAMES),
+}
+
+BOILERPLATE_PHRASES: List[str] = [
+    "provided proper attribution", "google hereby grants permission",
+    "permission to reproduce", "all rights reserved", "copyright ©",
+    "under the terms of", "creative commons", "open access article",
+    "preprint server", "arxiv:", "doi:", "received:", "accepted:",
+    "published online", "correspondence to",
+]
+
+METRIC_ONLY_TERMS: Set[str] = {
+    "accuracy", "acc", "precision", "recall", "sensitivity", "specificity",
+    "f1", "f1 score", "auc", "roc auc", "auroc", "auprc", "bleu",
+    "rouge", "rouge-l", "meteor", "map", "ndcg", "wer", "cer",
+    "perplexity", "ppl", "loss", "rmse", "mae", "mse", "iou",
+    "dice", "ari", "nmi", "top-1", "top-5", "er", "hr",
+}
+
+TABLE_CELL_HINTS: List[str] = [
+    "baseline", "ours", "oracle", "ensemble", "single model", "dev", "test",
+    "train", "training", "validation", "params", "flops", "gpu", "cpu",
+    "memory", "latency",
+]
+
+_HEADING_CONFIDENCE_THRESHOLD = 0.55
+_MIN_SECTION_WORDS = 20
+
+
+# ---------------------------------------------------------------------------
+# Text normalization
+# ---------------------------------------------------------------------------
+
+def _clean_text(text: str) -> str:
+    if not text:
+        return ""
+    text = (
+        text.replace("\x00", " ")
+        .replace("\u00a0", " ")
+        .replace("\u000f", " ")
+        .replace("\ufb01", "fi")
+        .replace("\ufb02", "fl")
+        .replace("\u2013", "-")
+        .replace("\u2014", "-")
+        .replace("\u2018", "'")
+        .replace("\u2019", "'")
+        .replace("\u201c", '"')
+        .replace("\u201d", '"')
+    )
+    text = re.sub(r"-\n(?=[a-z])", "", text)
+    text = re.sub(r"[ \t]+", " ", text)
+    text = re.sub(r"\n[ \t]+", "\n", text)
+    text = re.sub(r"\n{3,}", "\n\n", text)
+    return text.strip()
+
+
+def _normalize_heading(text: str) -> str:
+    text = _clean_text(text).strip()
+    text = re.sub(r"^[#*\s]+", "", text)
+    text = re.sub(r"^(section|chapter|part)\s+", "", text, flags=re.IGNORECASE)
+    text = re.sub(r"^[IVXLCDM]+[.)]\s+", "", text, flags=re.IGNORECASE)
+    text = re.sub(r"^[A-Z][.)]\s+", "", text)
+    text = re.sub(r"^\d+(?:\.\d+)*[.)]?\s+", "", text)
+    text = re.sub(r"[:.\s]+$", "", text)
+    return re.sub(r"\s+", " ", text).lower().strip()
+
+
+def _heading_level(text: str) -> int:
+    match = re.match(r"^(\d+(?:\.\d+)*)", text.strip())
+    if match:
+        return min(6, match.group(1).count(".") + 1)
+    return 1
+
+
+def _role_for_title(title: str) -> str:
+    norm = _normalize_heading(title)
+    for role, aliases in ROLE_ALIASES.items():
+        if norm in aliases:
+            return role
+    for role, aliases in ROLE_ALIASES.items():
+        if any(alias in norm for alias in aliases):
+            return role
+    return "other"
+
+
+def _line_key(text: str) -> str:
+    return re.sub(r"[^a-z0-9]+", " ", _normalize_heading(text)).strip()
+
+
+# ---------------------------------------------------------------------------
+# Low-level PyMuPDF line extraction
+# ---------------------------------------------------------------------------
+
+def _span_is_bold(span: Dict[str, Any]) -> bool:
+    font = span.get("font", "") or ""
+    flags = int(span.get("flags", 0) or 0)
+    return "bold" in font.lower() or "heavy" in font.lower() or bool(flags & 16)
+
+
+def _extract_page_lines(page: fitz.Page, page_number: int) -> List[Dict[str, Any]]:
+    data = page.get_text("dict", flags=fitz.TEXT_PRESERVE_WHITESPACE)
+    page_width = float(page.rect.width)
+    mid_x = page_width / 2.0
+    raw: List[Dict[str, Any]] = []
+
+    for block in data.get("blocks", []):
+        if block.get("type") != 0:
+            continue
+        block_no = block.get("number", 0)
+        for line in block.get("lines", []):
+            spans = [s for s in line.get("spans", []) if (s.get("text") or "").strip()]
+            if not spans:
+                continue
+            text = _clean_text("".join(s.get("text", "") for s in spans))
+            if not text:
+                continue
+
+            sizes = [round(float(s.get("size", 0.0)), 1) for s in spans if s.get("size")]
+            max_size = max(sizes) if sizes else 0.0
+            avg_size = sum(sizes) / len(sizes) if sizes else 0.0
+            bold_ratio = sum(1 for s in spans if _span_is_bold(s)) / max(1, len(spans))
+            bbox = tuple(line.get("bbox") or block.get("bbox") or (0, 0, 0, 0))
+            x0 = float(bbox[0])
+            y0 = float(bbox[1])
+            col = 0 if x0 < mid_x else 1
+
+            raw.append({
+                "text": text,
+                "page_number": page_number,
+                "block_no": block_no,
+                "bbox": bbox,
+                "x0": x0,
+                "y0": y0,
+                "max_size": max_size,
+                "avg_size": avg_size,
+                "bold": bold_ratio >= 0.5,
+                "span_count": len(spans),
+                "col": col,
+            })
+
+    if raw:
+        max_y = max(float(l["y0"]) for l in raw)
+        band_height = max(max_y / 20.0, 1.0)
+        raw.sort(key=lambda l: (round(float(l["y0"]) / band_height), int(l["col"]), float(l["y0"])))
+
+    return raw
+
+
+def _dominant_body_size(lines: List[Dict[str, Any]]) -> float:
+    sizes: List[float] = []
+    for line in lines:
+        size = float(line.get("avg_size") or line.get("max_size") or 0.0)
+        words = len(str(line.get("text", "")).split())
+        if 7.0 <= size <= 16.0 and words >= 4:
+            sizes.append(round(size, 1))
+    return Counter(sizes).most_common(1)[0][0] if sizes else 10.0
+
+
+# ---------------------------------------------------------------------------
+# Heading scoring
+# ---------------------------------------------------------------------------
+
+def _boilerplate_line(text: str) -> bool:
+    low = _clean_text(text).lower()
+    return any(phrase in low for phrase in BOILERPLATE_PHRASES)
+
+
+def _metric_only_heading(text: str) -> bool:
+    norm = _normalize_heading(text)
+    compact = re.sub(r"[^a-z0-9]+", " ", norm).strip()
+    if compact in METRIC_ONLY_TERMS:
+        return True
+
+    metric_re = r"\b(" + "|".join(re.escape(m) for m in sorted(METRIC_ONLY_TERMS, key=len, reverse=True)) + r")\b"
+    if len(text.split()) <= 5 and re.search(metric_re, compact, flags=re.IGNORECASE):
+        if not any(w in compact for w in ["results", "evaluation", "experiment", "analysis", "method"]):
+            return True
+
+    letters = re.sub(r"[^A-Za-z]", "", text)
+    if 2 <= len(letters) <= 10 and letters.isupper() and len(text.split()) <= 3:
+        return True
+    return False
+
+
+def _probable_table_cell(text: str) -> bool:
+    raw = _clean_text(text)
+    low = raw.lower()
+    words = raw.split()
+    if not raw or len(words) > 9:
+        return False
+
+    if re.search(r"\[\d+\]", raw) and len(words) <= 7:
+        return True
+
+    chars = [c for c in raw if not c.isspace()]
+    if chars:
+        symbol_digit_ratio = sum(1 for c in chars if c.isdigit() or c in ".,±+-×*/=()%[]") / len(chars)
+        if symbol_digit_ratio >= 0.35 and len(words) <= 6:
+            return True
+
+    if len(words) <= 5 and any(h in low for h in TABLE_CELL_HINTS):
+        if not any(sw in low for sw in ["result", "method", "experiment", "evaluation", "discussion", "conclusion"]):
+            return True
+
+    if len(words) <= 5 and re.search(r"\b[A-Z][A-Za-z0-9-]*\s*(\+|/|vs\.?|&|×)\s*[A-Z]", raw):
+        return True
+
+    return False
+
+
+def _looks_like_reference_item(line: str) -> bool:
+    return bool(re.match(r"^(\[\d+\]|\d+[.)])\s+", line.strip())) and len(line.split()) > 6
+
+
+_NUMBERED_HEADING_RE = re.compile(
+    r"^(?:\d+(?:\.\d+)*[.)]?|[IVXLCDM]+[.)]|[A-Z][.)])\s+"
+    r"[A-Z\u00C0-\u024F][A-Za-z0-9\u00C0-\u024F ,/()\-–—:&']+$",
+    re.IGNORECASE,
+)
+
+
+def _heading_confidence(
+    line: Dict[str, Any],
+    body_size: float,
+    prev_line: Optional[Dict[str, Any]],
+    next_line: Optional[Dict[str, Any]],
+) -> float:
+    text = _clean_text(line.get("text", ""))
+    if not text:
+        return 0.0
+    if _boilerplate_line(text) or _metric_only_heading(text) or _probable_table_cell(text):
+        return 0.0
+    if _looks_like_reference_item(text):
+        return 0.0
+    if len(text) > 150 or len(text.split()) > 18:
+        return 0.0
+    if text.endswith(",") or text.endswith(";"):
+        return 0.0
+
+    norm = _normalize_heading(text)
+    if not norm or norm in {"figure", "table", "fig", "eq", "equation"}:
+        return 0.0
+
+    score = 0.0
+
+    if norm in SECTION_KEYWORDS:
+        score += 0.70
+    else:
+        for keyword in SECTION_KEYWORDS:
+            if keyword in norm and len(keyword) >= 6:
+                score += 0.30
+                break
+
+    if _NUMBERED_HEADING_RE.match(text) and len(text.split()) <= 14:
+        score += 0.55
+
+    letters = re.sub(r"[^A-Za-z]", "", text)
+    if len(letters) >= 4 and letters.isupper() and len(text.split()) <= 10:
+        score += 0.40
+
+    size = float(line.get("max_size") or line.get("avg_size") or body_size)
+    size_ratio = size / body_size if body_size else 1.0
+    if size_ratio >= 1.25:
+        score += 0.35
+    elif size_ratio >= 1.10:
+        score += 0.20
+
+    if line.get("bold"):
+        score += 0.20
+
+    bbox = line.get("bbox") or (0, 0, 0, 0)
+    y0 = float(bbox[1])
+    y1 = float(bbox[3])
+    prev_gap = 999.0
+    next_gap = 999.0
+    if prev_line and prev_line.get("page_number") == line.get("page_number"):
+        prev_gap = y0 - float((prev_line.get("bbox") or (0, 0, 0, 0))[3])
+    if next_line and next_line.get("page_number") == line.get("page_number"):
+        next_gap = float((next_line.get("bbox") or (0, 0, 0, 0))[1]) - y1
+    if prev_gap >= 3.0 or next_gap >= 2.0:
+        score += 0.15
+
+    if re.match(r"^[A-Z\u00C0-\u024F][A-Za-z0-9\u00C0-\u024F ,/()\-–—:&']+$", text):
+        if not re.search(r"\b(the|and|or|but|because|while|which|that)\b.+\.$", text.lower()):
+            score += 0.10
+
+    if len(text.split()) < 2 and score < 0.70:
+        score *= 0.40
+
+    return min(score, 1.0)
+
+
+def _detect_heading_keys(lines: List[Dict[str, Any]], body_size: float) -> Set[str]:
+    candidate_scores: Dict[str, List[float]] = defaultdict(list)
+
+    for i, line in enumerate(lines):
+        prev_line = lines[i - 1] if i else None
+        next_line = lines[i + 1] if i + 1 < len(lines) else None
+        confidence = _heading_confidence(line, body_size, prev_line, next_line)
+        if confidence > 0.0:
+            key = _line_key(line["text"])
+            if key:
+                candidate_scores[key].append(confidence)
+
+    accepted: Set[str] = set()
+    for key, scores in candidate_scores.items():
+        best = max(scores)
+        if best >= _HEADING_CONFIDENCE_THRESHOLD:
+            accepted.add(key)
+        elif best >= 0.35 and len(scores) >= 2:
+            accepted.add(key)
+    return accepted
+
+
+# ---------------------------------------------------------------------------
+# Section splitting and back-matter separation
+# ---------------------------------------------------------------------------
+
+def _split_sections_from_lines(lines: List[Dict[str, Any]], heading_keys: Set[str]) -> List[Dict[str, Any]]:
+    sections: List[Dict[str, Any]] = []
+    current_title = "Front Matter"
+    current_level = 1
+    current_role = "front_matter"
+    current_lines: List[str] = []
+    page_start: Optional[int] = lines[0]["page_number"] if lines else None
+    page_end: Optional[int] = page_start
+
+    def flush() -> None:
+        nonlocal current_lines, current_title, current_level, current_role, page_start, page_end
+        text = _clean_text("\n".join(current_lines))
+        if text:
+            sections.append({
+                "title": current_title,
+                "text": text,
+                "level": current_level,
+                "role": current_role,
+                "page_start": page_start,
+                "page_end": page_end,
+                "word_count": len(text.split()),
+            })
+        current_lines = []
+
+    for line in lines:
+        text = _clean_text(line["text"])
+        key = _line_key(text)
+        if key in heading_keys:
+            flush()
+            current_title = text
+            current_level = _heading_level(text)
+            current_role = _role_for_title(text)
+            page_start = line.get("page_number")
+            page_end = page_start
+        else:
+            current_lines.append(text)
+            page_end = line.get("page_number")
+
+    flush()
+
+    merged: List[Dict[str, Any]] = []
+    for sec in sections:
+        norm = _normalize_heading(sec["title"])
+        words = int(sec.get("word_count", len(sec.get("text", "").split())))
+        if sec["title"] != "Front Matter" and norm not in SECTION_KEYWORDS and words < _MIN_SECTION_WORDS and merged:
+            merged[-1]["text"] = _clean_text(merged[-1]["text"] + "\n" + sec["title"] + "\n" + sec["text"])
+            merged[-1]["page_end"] = sec.get("page_end") or merged[-1].get("page_end")
+            merged[-1]["word_count"] = len(merged[-1]["text"].split())
+        else:
+            merged.append(sec)
+    return merged
+
+
+def _separate_back_matter(sections: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], str, str, str]:
+    body: List[Dict[str, Any]] = []
+    references_parts: List[str] = []
+    appendix_parts: List[str] = []
+    boilerplate_parts: List[str] = []
+
+    mode = "body"
+    for sec in sections:
+        role = sec.get("role") or _role_for_title(sec.get("title", ""))
+        norm = _normalize_heading(sec.get("title", ""))
+        packed = _clean_text(f'{sec.get("title", "")}\n{sec.get("text", "")}')
+
+        if role == "references" or norm in REFERENCE_SECTION_NAMES:
+            mode = "references"
+        elif role == "appendix" or norm in APPENDIX_SECTION_NAMES:
+            if mode != "references":
+                mode = "appendix"
+        elif role == "boilerplate" or norm in BOILERPLATE_SECTION_NAMES:
+            if mode not in {"references", "appendix"}:
+                mode = "boilerplate"
+        elif mode == "boilerplate":
+            mode = "body"
+
+        if mode == "references":
+            references_parts.append(packed)
+        elif mode == "appendix":
+            appendix_parts.append(packed)
+        elif mode == "boilerplate":
+            boilerplate_parts.append(packed)
+        else:
+            body.append(sec)
+
+    return body, "\n\n".join(references_parts), "\n\n".join(appendix_parts), "\n\n".join(boilerplate_parts)
+
+
+# ---------------------------------------------------------------------------
+# Title, abstract, references, captions, tables
+# ---------------------------------------------------------------------------
+
+_BAD_TITLE_MARKERS = [
+    "provided proper attribution", "google hereby grants permission", "arxiv",
+    "preprint", "license", "copyright", "all rights reserved",
+]
+
+_BLOCKED_TITLE_FRAGMENTS = [
+    "vol.", "http", "www.", "received", "accepted", "department", "university",
+    "hospital", "college", "institute", "email", "@", "proceedings",
+    "journal of", "conference on",
+]
+
+
+def _extract_title(lines: List[Dict[str, Any]]) -> Optional[str]:
+    first_page = [line for line in lines if line.get("page_number") == 1]
+    candidates: List[Dict[str, Any]] = []
+
+    for line in first_page[:100]:
+        text = _clean_text(line["text"])
+        low = text.lower()
+        if len(text) < 8 or len(text) > 200:
+            continue
+        if any(x in low for x in _BAD_TITLE_MARKERS + _BLOCKED_TITLE_FRAGMENTS):
+            continue
+        if re.match(r"^\d+$", text) or _normalize_heading(text) in SECTION_KEYWORDS:
+            continue
+        if text.endswith(".") and len(text.split()) < 10:
+            continue
+        candidates.append(line)
+
+    if not candidates:
+        return None
+
+    max_size = max(float(c.get("max_size") or 0) for c in candidates)
+    title_lines = [c for c in candidates if float(c.get("max_size") or 0) >= max_size - 0.6]
+
+    by_block: Dict[int, List[str]] = defaultdict(list)
+    for line in title_lines:
+        by_block[int(line.get("block_no", 0))].append(line["text"])
+
+    for _, parts in sorted(by_block.items()):
+        title = _clean_text(" ".join(parts))
+        if 3 <= len(title.split()) <= 35:
+            return title
+
+    for line in candidates[:25]:
+        text = _clean_text(line["text"])
+        if 3 <= len(text.split()) <= 30 and not text.endswith("."):
+            return text
+    return None
+
+
+_ABSTRACT_MARKERS = [
+    "we propose", "we present", "we introduce", "we developed", "we investigate",
+    "this paper", "this work", "this study", "this article", "here we",
+    "in this paper", "in this work", "in this study",
+]
+
+
+def _extract_abstract_from_sections(sections: List[Dict[str, Any]], clean_text: str) -> Optional[str]:
+    for sec in sections:
+        if sec.get("role") == "abstract" or _normalize_heading(sec.get("title", "")) == "abstract":
+            text = _clean_text(sec.get("text", ""))
+            text = re.sub(r"^abstract[:\s—–-]*", "", text, flags=re.IGNORECASE).strip()
+            if len(text.split()) >= 30:
+                return text
+
+    explicit = re.search(
+        r"\babstract\b[:\s—–-]*(.{100,2500}?)(?=\n\s*(?:\d+[.)]?\s*)?(?:keywords|introduction|background|1[.\s]|i[.\s])\b)",
+        clean_text,
+        flags=re.IGNORECASE | re.DOTALL,
+    )
+    if explicit:
+        candidate = _clean_text(explicit.group(1))
+        if len(candidate.split()) >= 30:
+            return candidate
+
+    paragraphs = [_clean_text(p) for p in re.split(r"\n\s*\n", clean_text[:8000]) if _clean_text(p)]
+    for para in paragraphs[:10]:
+        low = para.lower()
+        if len(para.split()) >= 40 and any(marker in low for marker in _ABSTRACT_MARKERS):
+            return para
+    return None
+
+
+_REF_SPLIT_RE = re.compile(
+    r"\n(?=\s*(?:\[\d+\]|\d+[.)]|[A-Z][a-zA-Z\-]+,\s+[A-Z]|\(\d{4}\)))"
+)
+
+
+def _parse_references(references_text: str) -> List[str]:
+    references_text = _clean_text(references_text)
+    if not references_text:
+        return []
+    references_text = re.sub(r"^references\s*", "", references_text, flags=re.IGNORECASE).strip()
+    parts = _REF_SPLIT_RE.split(references_text)
+
+    refs: List[str] = []
+    for part in parts:
+        ref = _clean_text(part)
+        if len(ref) >= 25 and _normalize_heading(ref) not in BOILERPLATE_SECTION_NAMES:
+            refs.append(ref)
+
+    seen: Set[str] = set()
+    unique: List[str] = []
+    for ref in refs:
+        key = re.sub(r"\s+", " ", ref.lower())[:160]
+        if key not in seen:
+            seen.add(key)
+            unique.append(ref)
+    return unique[:300]
+
+
+_CAPTION_RE = re.compile(
+    r"^(?P<label>"
+    r"(?:figure|fig\.?|table|tbl\.?|scheme|supplementary figure|supp\.?\s*fig\.?|"
+    r"extended data figure|algorithm|listing)\s*[\dIVXLC]+[A-Za-z]?"
+    r")[\.:：\s\-–—]*(?P<caption>.+)$",
+    re.IGNORECASE,
+)
+
+
+def _extract_captions(lines: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    captions: List[Dict[str, Any]] = []
+    seen: Set[str] = set()
+    texts = [line["text"] for line in lines]
+
+    for i, line in enumerate(lines):
+        m = _CAPTION_RE.match(_clean_text(line["text"]))
+        if not m:
+            continue
+        label = m.group("label").strip()
+        caption = m.group("caption").strip()
+        if re.match(r"^(shows|illustrates|presents|depicts|demonstrates)\b", caption.lower()):
+            continue
+
+        continuation: List[str] = []
+        for next_text in texts[i + 1:i + 6]:
+            nt = _clean_text(next_text)
+            if _CAPTION_RE.match(nt) or _normalize_heading(nt) in SECTION_KEYWORDS:
+                break
+            if len(nt.split()) < 4:
+                break
+            continuation.append(nt)
+        if continuation:
+            caption = _clean_text(caption + " " + " ".join(continuation))
+
+        key = re.sub(r"[^a-z0-9]+", "", label.lower())
+        if key in seen:
+            continue
+        seen.add(key)
+        captions.append({
+            "label": label,
+            "caption": caption,
+            "page_number": line.get("page_number"),
+        })
+    return captions
+
+
+def _is_valid_table(data: Any) -> bool:
+    if not isinstance(data, list) or len(data) < 2:
+        return False
+    rows = [row for row in data if isinstance(row, list)]
+    if len(rows) < 2:
+        return False
+    cols = max((len(row) for row in rows), default=0)
+    if cols < 2:
+        return False
+
+    cells = [str(cell).strip() for row in rows for cell in row if cell is not None and str(cell).strip()]
+    non_empty = len(cells)
+    total_cells = sum(len(row) for row in rows)
+    if non_empty < 6:
+        return False
+
+    flat = " ".join(cells).lower()
+    if non_empty <= 8 and any(noise in flat for noise in ["www.", "http", "vol.", "doi"]):
+        return False
+
+    fill_ratio = non_empty / max(total_cells, 1)
+    avg_cell_len = sum(len(cell) for cell in cells) / max(non_empty, 1)
+    one_char_ratio = sum(1 for cell in cells if len(cell) <= 1) / max(non_empty, 1)
+    very_short_ratio = sum(1 for cell in cells if len(cell) <= 2) / max(non_empty, 1)
+
+    if cols > 12 and (fill_ratio < 0.35 or avg_cell_len < 3.0 or one_char_ratio > 0.35):
+        return False
+    if avg_cell_len < 2.5 and very_short_ratio > 0.60:
+        return False
+    if cols > 20 and non_empty < 80:
+        return False
+
+    meaningful_rows = sum(
+        1 for row in rows
+        if len([str(c).strip() for c in row if c is not None and str(c).strip()]) >= 2
+        and len(" ".join(str(c) for c in row)) >= 12
+    )
+    return meaningful_rows >= 2
+
+
+def _quality_report(result: Dict[str, Any]) -> Dict[str, Any]:
+    title = result.get("title") or ""
+    abstract = result.get("abstract") or ""
+    sections = result.get("sections", [])
+    roles = {section.get("role", "other") for section in sections}
+    n_refs = len(result.get("references", []))
+
+    score = 0.0
+    score += 0.15 if len(title.split()) >= 4 else 0.0
+    score += 0.20 if len(abstract.split()) >= 30 else 0.0
+    score += min(0.20, len(sections) * 0.022)
+    score += 0.10 if "methodology" in roles else 0.0
+    score += 0.08 if {"results", "experiments"} & roles else 0.0
+    score += 0.07 if "conclusion" in roles else 0.0
+    score += 0.08 if n_refs >= 5 else (0.03 if n_refs else 0.0)
+    score += 0.05 if len(result.get("clean_text", "")) > 4000 else 0.0
+    score += 0.04 if result.get("captions") else 0.0
+    score += 0.03 if result.get("tables") else 0.0
+
+    if len(abstract.split()) < 20:
+        score -= 0.10
+    if not sections:
+        score -= 0.20
+
+    return {
+        "title_found": len(title.split()) >= 4,
+        "abstract_found": len(abstract.split()) >= 30,
+        "num_sections": len(sections),
+        "section_roles": sorted(roles),
+        "methodology_section_found": "methodology" in roles,
+        "num_references": n_refs,
+        "num_captions": len(result.get("captions", [])),
+        "num_tables": len(result.get("tables", [])),
+        "references_removed_from_clean_text": bool(result.get("references_text")),
+        "appendix_removed_from_clean_text": bool(result.get("appendix_text")),
+        "quality_score": round(max(0.0, min(score, 0.98)), 3),
+    }
+
+
+# ---------------------------------------------------------------------------
+# Engines
+# ---------------------------------------------------------------------------
+
+def extract_with_pymupdf(pdf_path: str | Path) -> Dict[str, Any]:
+    path = Path(pdf_path)
+    if not path.exists():
+        raise PDFIngestionError(f"PDF not found: {path}")
+
+    doc = fitz.open(path)
+    try:
+        all_lines: List[Dict[str, Any]] = []
+        pages: List[Dict[str, Any]] = []
+        tables: List[Dict[str, Any]] = []
+
+        for page_index, page in enumerate(doc):
+            page_number = page_index + 1
+            page_lines = _extract_page_lines(page, page_number)
+            all_lines.extend(page_lines)
+            page_text = _clean_text("\n".join(line["text"] for line in page_lines))
+            pages.append({"page_number": page_number, "text": page_text})
+
+            try:
+                found = page.find_tables()
+                for table_index, table in enumerate(found.tables):
+                    data = table.extract()
+                    if _is_valid_table(data):
+                        tables.append({
+                            "page_number": page_number,
+                            "table_index": table_index,
+                            "data": data,
+                            "engine": "pymupdf",
+                            "caption": None,
+                        })
+            except Exception:
+                pass
+
+        full_text = _clean_text("\n\n".join(page["text"] for page in pages))
+        body_size = _dominant_body_size(all_lines)
+        heading_keys = _detect_heading_keys(all_lines, body_size)
+        sections_all = _split_sections_from_lines(all_lines, heading_keys)
+        body_sections, references_text, appendix_text, boilerplate_text = _separate_back_matter(sections_all)
+        clean_text = _clean_text("\n\n".join(f'{s["title"]}\n{s["text"]}' for s in body_sections))
+
+        title = _extract_title(all_lines)
+        abstract = _extract_abstract_from_sections(body_sections, clean_text)
+        references = _parse_references(references_text)
+        captions = _extract_captions(all_lines)
+
+        # Attach table captions when labels and table indices align.
+        cap_map: Dict[str, str] = {}
+        for cap in captions:
+            label = cap.get("label", "").lower()
+            if "table" in label or "tbl" in label:
+                key = re.sub(r"[^a-z0-9]", "", label)
+                cap_map[key] = cap.get("caption", "")
+        for table in tables:
+            idx_str = str(table.get("table_index", ""))
+            for key, caption in cap_map.items():
+                if idx_str in key or key.endswith(idx_str):
+                    table["caption"] = caption
+                    break
+
+        result: Dict[str, Any] = {
+            "source_pdf": path.name,
+            "num_pages": len(doc),
+            "title": title,
+            "abstract": abstract,
+            "text": clean_text,
+            "clean_text": clean_text,
+            "raw_text": full_text,
+            "pages": pages,
+            "sections": body_sections,
+            "all_sections": sections_all,
+            "references": references,
+            "references_text": references_text,
+            "appendix_text": appendix_text,
+            "boilerplate_text": boilerplate_text,
+            "captions": captions,
+            "tables": tables,
+            "metadata": {
+                "body_font_size": body_size,
+                "heading_count": len(heading_keys),
+                "removed_back_matter": {
+                    "references": bool(references_text),
+                    "appendix": bool(appendix_text),
+                    "boilerplate": bool(boilerplate_text),
+                },
+            },
+            "extraction_engine": "pymupdf-section-aware-final",
+        }
+        result["quality"] = _quality_report(result)
+        return result
+    finally:
+        doc.close()
+
+
+def extract_with_docling(pdf_path: str | Path) -> Dict[str, Any]:
+    path = Path(pdf_path)
+    if not path.exists():
+        raise PDFIngestionError(f"PDF not found: {path}")
+
+    try:
+        from docling.document_converter import DocumentConverter
+    except ImportError as exc:
+        raise PDFIngestionError("Docling is not installed. Run: pip install docling") from exc
+
+    converter = DocumentConverter()
+    doc = converter.convert(str(path)).document
+    markdown = _clean_text(doc.export_to_markdown())
+
+    fake_lines: List[Dict[str, Any]] = []
+    for i, raw in enumerate(markdown.splitlines()):
+        is_heading = raw.startswith("#")
+        text = _clean_text(re.sub(r"^#+\s*", "", raw))
+        if text:
+            fake_lines.append({
+                "text": text,
+                "page_number": None,
+                "block_no": i,
+                "bbox": (0, i * 10, 0, i * 10 + 8),
+                "max_size": 13.0 if is_heading else 10.0,
+                "avg_size": 13.0 if is_heading else 10.0,
+                "bold": is_heading,
+                "col": 0,
+            })
+
+    heading_keys = _detect_heading_keys(fake_lines, 10.0)
+    for line in fake_lines:
+        if line.get("bold") and _normalize_heading(line["text"]) in SECTION_KEYWORDS:
+            heading_keys.add(_line_key(line["text"]))
+
+    sections_all = _split_sections_from_lines(fake_lines, heading_keys)
+    body_sections, references_text, appendix_text, boilerplate_text = _separate_back_matter(sections_all)
+    clean_text = _clean_text("\n\n".join(f'{s["title"]}\n{s["text"]}' for s in body_sections))
+
+    result: Dict[str, Any] = {
+        "source_pdf": path.name,
+        "num_pages": None,
+        "title": _extract_title(fake_lines),
+        "abstract": _extract_abstract_from_sections(body_sections, clean_text),
+        "text": clean_text,
+        "clean_text": clean_text,
+        "raw_text": markdown,
+        "pages": [],
+        "sections": body_sections,
+        "all_sections": sections_all,
+        "references": _parse_references(references_text),
+        "references_text": references_text,
+        "appendix_text": appendix_text,
+        "boilerplate_text": boilerplate_text,
+        "captions": [],
+        "tables": [],
+        "metadata": {
+            "body_font_size": 10.0,
+            "heading_count": len(heading_keys),
+            "removed_back_matter": {
+                "references": bool(references_text),
+                "appendix": bool(appendix_text),
+                "boilerplate": bool(boilerplate_text),
+            },
+        },
+        "extraction_engine": "docling-section-aware-final",
+    }
+    result["quality"] = _quality_report(result)
+    return result
+
+
+def extract_pdf(pdf_path: str | Path, engine: str = "pymupdf") -> Dict[str, Any]:
+    engine = engine.lower().strip()
+    if engine == "pymupdf":
+        return extract_with_pymupdf(pdf_path)
+    if engine == "docling":
+        return extract_with_docling(pdf_path)
+    if engine == "auto":
+        try:
+            pymupdf_result = extract_with_pymupdf(pdf_path)
+            try:
+                docling_result = extract_with_docling(pdf_path)
+                return (
+                    pymupdf_result
+                    if pymupdf_result["quality"]["quality_score"] >= docling_result["quality"]["quality_score"]
+                    else docling_result
+                )
+            except PDFIngestionError:
+                return pymupdf_result
+        except Exception:
+            return extract_with_docling(pdf_path)
+    raise PDFIngestionError(f"Unknown engine '{engine}'. Choose: pymupdf | docling | auto.")

src/paper2lab/data/schema.py

@@ -0,0 +1,92 @@
+"""
+schema.py — Paper2Lab data contracts.
+
+Field-agnostic contracts for PDF extraction across ML, NLP, CV, biomedical,
+physics, education, social-science, economics, and interdisciplinary papers.
+"""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from typing import Any, Dict, List, Optional
+
+
+KNOWN_ROLES = frozenset({
+    "front_matter",
+    "abstract",
+    "keywords",
+    "introduction",
+    "related_work",
+    "background",
+    "theory",
+    "methodology",
+    "experiments",
+    "results",
+    "discussion",
+    "limitations",
+    "future_work",
+    "conclusion",
+    "references",
+    "appendix",
+    "boilerplate",
+    "other",
+})
+
+
+@dataclass
+class Section:
+    title: str
+    text: str
+    level: int = 1
+    page_start: Optional[int] = None
+    page_end: Optional[int] = None
+    role: str = "other"
+    word_count: int = 0
+
+    def __post_init__(self) -> None:
+        if self.role not in KNOWN_ROLES:
+            self.role = "other"
+        if not self.word_count:
+            self.word_count = len((self.text or "").split())
+
+
+@dataclass
+class Caption:
+    label: str
+    caption: str
+    page_number: Optional[int] = None
+
+
+@dataclass
+class Table:
+    page_number: Optional[int]
+    table_index: int
+    data: Any
+    engine: str = "unknown"
+    caption: Optional[str] = None
+
+
+@dataclass
+class DocumentExtraction:
+    source_pdf: str
+    title: Optional[str]
+    abstract: Optional[str]
+    text: str
+    clean_text: str
+    raw_text: str = ""
+    num_pages: Optional[int] = None
+    sections: List[Section] = field(default_factory=list)
+    all_sections: List[Section] = field(default_factory=list)
+    references: List[str] = field(default_factory=list)
+    references_text: str = ""
+    appendix_text: str = ""
+    boilerplate_text: str = ""
+    captions: List[Caption] = field(default_factory=list)
+    tables: List[Table] = field(default_factory=list)
+    pages: List[Dict[str, Any]] = field(default_factory=list)
+    metadata: Dict[str, Any] = field(default_factory=dict)
+    quality: Dict[str, Any] = field(default_factory=dict)
+    extraction_engine: str = "unknown"
+
+    def to_dict(self) -> Dict[str, Any]:
+        return asdict(self)