paper_side.py

#!/usr/bin/env python3
"""
The paper side of a training sample: which assets exist, and what each is called.

Two things here are not what the earlier docling-based version did.

**Assets come from the LayoutLMv3 registry.** The poster side was derendered
with LayoutLMv3, so the paper side has to use it too or the crops never line up
-- that asymmetry is what made the first matching attempt useless. The registry
also carries page and bbox, which the caption logic needs.

**Captions are read out of the PDF at the caption box's own coordinates.**
Pairing the Nth crop with the Nth caption in a parsed-content dump is what put
"Figure 3"'s text on figure_4 earlier: any dropped detection on either side
shifts every later pairing silently. LayoutLMv3 already emits figure_caption /
table_caption boxes, so instead of counting, this finds the caption box nearest
*below* each asset on the same page and pulls the text inside that rectangle
straight from the PDF. A caption either has a box under its figure or the asset
goes uncaptioned; it can no longer be given someone else's.

    from paper_side import load_paper_assets
    assets = load_paper_assets(pid, paper_assets_root, pdf_path)
"""

from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Dict, List, Optional

CAPTION_FOR = {'figure': 'figure_caption', 'table': 'table_caption',
               'isolate_formula': 'formula_caption'}
KIND_NAME = {'figure': 'figure', 'table': 'table', 'isolate_formula': 'formula'}
# A caption sits directly under (tables: often over) its asset. Anything further
# than this fraction of the page height away belongs to something else.
MAX_CAPTION_GAP = 0.12
MAX_CAPTION_CHARS = 300


def _overlap(a, b) -> float:
    """Horizontal overlap of two boxes as a fraction of the narrower one."""
    lo, hi = max(a[0], b[0]), min(a[2], b[2])
    if hi <= lo:
        return 0.0
    return (hi - lo) / max(min(a[2] - a[0], b[2] - b[0]), 1e-6)


def _clean(t: str) -> str:
    t = re.sub(r'\s+', ' ', t).strip()
    if len(t) > MAX_CAPTION_CHARS:
        t = t[:MAX_CAPTION_CHARS].rsplit(' ', 1)[0] + '…'
    return t


def load_paper_assets(pid: str, root: str, pdf_path: Optional[str] = None,
                      main_body_pages: int = 12) -> Optional[Dict]:
    """-> {'assets': [...], 'counts': {...}, 'n_pages': int} or None."""
    d = Path(root) / pid
    meta_p = d / 'assets.json'
    if not meta_p.exists():
        return None
    try:
        meta = json.loads(meta_p.read_text(encoding='utf-8'))
    except Exception:                                                # noqa: BLE001
        return None

    assets = []
    for a in meta.get('assets', []):
        f = d / 'assets' / f"{a['id']}.png"
        if not f.exists():
            continue
        assets.append({'id': a['id'], 'class': a['class'], 'kind': KIND_NAME[a['class']],
                       'page': a['page'], 'bbox': a['bbox'], 'aspect': a.get('aspect'),
                       'area': (a['bbox'][2] - a['bbox'][0]) * (a['bbox'][3] - a['bbox'][1]),
                       'file': str(f), 'caption': '',
                       'main_body': a['page'] <= main_body_pages})

    if pdf_path:
        try:
            _attach_captions(assets, d, pdf_path, meta.get('dpi', 200))
        except Exception:                                            # noqa: BLE001
            pass                       # captions are a nicety, assets are not

    return {'paper_id': pid, 'assets': assets, 'counts': meta.get('counts', {}),
            'n_pages': meta.get('n_pages', 0)}


def _attach_captions(assets: List[Dict], paper_dir: Path, pdf_path: str,
                     dpi: int) -> None:
    layout_p = paper_dir / 'layout.json'
    if not layout_p.exists():
        return
    layout = json.loads(layout_p.read_text(encoding='utf-8'))
    caps: Dict[int, List[Dict]] = {}
    for d in layout:
        if d.get('class_name') in CAPTION_FOR.values():
            caps.setdefault(d['page'], []).append(d)
    if not caps:
        return

    import fitz
    doc = fitz.open(pdf_path)
    scale = 72.0 / dpi              # rendered pixels back to PDF points
    try:
        for a in assets:
            want = CAPTION_FOR.get(a['class'])
            page_caps = [c for c in caps.get(a['page'], [])
                         if c.get('class_name') == want]
            if not page_caps or a['page'] > len(doc):
                continue
            page = doc[a['page'] - 1]
            ph = page.rect.height / scale       # page height in rendered pixels
            ax1, ay1, ax2, ay2 = a['bbox']
            best, best_gap = None, 1e9
            for c in page_caps:
                cx1, cy1, cx2, cy2 = c['bbox']
                if _overlap((ax1, ay1, ax2, ay2), (cx1, cy1, cx2, cy2)) < 0.3:
                    continue
                # below the asset, or above it (the usual place for a table)
                gap = cy1 - ay2 if cy1 >= ay2 else ay1 - cy2
                if gap < -5 or gap > MAX_CAPTION_GAP * ph:
                    continue
                if gap < best_gap:
                    best, best_gap = c, gap
            if best is None:
                continue
            x1, y1, x2, y2 = (v * scale for v in best['bbox'])
            txt = page.get_text('text', clip=fitz.Rect(x1 - 2, y1 - 2, x2 + 2, y2 + 2))
            a['caption'] = _clean(txt)
    finally:
        doc.close()


# --------------------------------------------------------------------------- #
# paper text
# --------------------------------------------------------------------------- #

# LayoutLMv3's own vocabulary. `abandon` is what it calls running heads, page
# numbers and footers -- exactly the text a poster never carries, and exactly
# what a raw page-level text dump would splice into the middle of a paragraph.
TEXT_CLASSES = {'title': 'title', 'plain text': 'text', 'plain_text': 'text'}
DROP_CLASSES = {'abandon', 'figure', 'table', 'isolate_formula',
                'figure_caption', 'table_caption', 'formula_caption',
                'table_footnote'}
MAX_BODY_CHARS = 12000


SPAN_FRAC = 0.75          # of the page's own text width, not the paper size


def is_two_column(boxes: List[Dict], mid: float) -> bool:
    """Does this page actually have two columns?

    Deciding by "is the centre left or right of the page middle" invents a
    gutter on single-column pages, and a wide-margin single-column layout is
    common enough (ICLR's template among them) that the invented gutter
    reordered abstracts after introductions. A real two-column page has boxes
    that finish before the middle *and* boxes that start after it.
    """
    return (any(b['bbox'][2] < mid for b in boxes)
            and any(b['bbox'][0] > mid for b in boxes))


def reading_order(boxes: List[Dict], page_width: float) -> List[Dict]:
    """Order text boxes the way the page is read.

    Two corrections over sorting by (column, y):

    Spanning is measured against the widest box on the page, not the paper
    size. A single-column ICLR page is 1700 px wide but its text column is
    1115 px, so a page-relative threshold classified ordinary body paragraphs
    as spanning and full-width titles as not.

    Column splitting only happens on pages that have two columns. Otherwise
    everything is read top to bottom, which is what a single-column page wants
    and what the gutter heuristic actively broke.

    On a genuine two-column page a spanning box acts as a barrier: pending
    column boxes are emitted left-then-right before it, so full-width headers
    stay in place and the columns beneath them stay in order.
    """
    if not boxes:
        return []
    body_w = max(b['bbox'][2] - b['bbox'][0] for b in boxes)
    span_w = SPAN_FRAC * body_w
    ordered = sorted(boxes, key=lambda b: b['bbox'][1])

    mid = page_width / 2
    narrow = [b for b in ordered if (b['bbox'][2] - b['bbox'][0]) <= span_w]
    if not is_two_column(narrow, mid):
        return ordered

    out: List[Dict] = []
    pending: List[Dict] = []

    def flush():
        if not pending:
            return
        left = [b for b in pending if (b['bbox'][0] + b['bbox'][2]) / 2 < mid]
        right = [b for b in pending if (b['bbox'][0] + b['bbox'][2]) / 2 >= mid]
        for group in (left, right):
            out.extend(sorted(group, key=lambda b: b['bbox'][1]))
        pending.clear()

    for b in ordered:
        if (b['bbox'][2] - b['bbox'][0]) > span_w:
            flush()
            out.append(b)
        else:
            pending.append(b)
    flush()
    return out


def load_paper_text(pid: str, root: str, pdf_path: str,
                    max_pages: int = 12, max_chars: int = MAX_BODY_CHARS) -> str:
    """Body text in reading order, taken from the detected text boxes.

    Reading text box by box rather than page by page buys three things the raw
    dump cannot: running heads and page numbers are dropped because LayoutLMv3
    labels them `abandon`; captions and equations are dropped because they are
    already carried as assets and would otherwise appear twice; and two-column
    papers come out in column order instead of interleaved across the gutter.

    Ordering is by column then vertical position, which is what a two-column
    layout needs and what a single-column layout is unaffected by.
    """
    d = Path(root) / pid
    layout_p = d / 'layout.json'
    if not layout_p.exists() or not pdf_path:
        return ''
    try:
        layout = json.loads(layout_p.read_text(encoding='utf-8'))
        meta = json.loads((d / 'assets.json').read_text(encoding='utf-8'))
    except Exception:                                                # noqa: BLE001
        return ''
    scale = 72.0 / meta.get('dpi', 200)

    try:
        import fitz
        doc = fitz.open(pdf_path)
    except Exception:                                                # noqa: BLE001
        return ''

    parts: List[str] = []
    try:
        for page_no in range(1, min(max_pages, len(doc)) + 1):
            boxes = [b for b in layout
                     if b.get('page') == page_no
                     and b.get('class_name') in TEXT_CLASSES
                     and b.get('class_name') not in DROP_CLASSES]
            if not boxes:
                continue
            page = doc[page_no - 1]
            pw = page.rect.width / scale
            for b in reading_order(boxes, pw):
                x1, y1, x2, y2 = (v * scale for v in b['bbox'])
                t = page.get_text('text', clip=fitz.Rect(x1 - 1, y1 - 1, x2 + 1, y2 + 1))
                t = re.sub(r'\s+', ' ', t).strip()
                if len(t) < 3:
                    continue
                parts.append(('## ' + t) if TEXT_CLASSES[b['class_name']] == 'title' else t)
                if sum(len(p) for p in parts) > max_chars:
                    raise StopIteration
    except StopIteration:
        pass
    finally:
        doc.close()

    out = '\n\n'.join(parts)
    return out[:max_chars]


def select_for_prompt(assets: List[Dict], required: List[str],
                      max_total: int, per_kind: Dict[str, int]) -> List[Dict]:
    """Choose which assets the model is shown.

    Anything the target program references is mandatory: dropping it would make
    the sample ask for a figure that was never offered, teaching the model to
    hallucinate ids. Whatever budget is left goes to the largest remaining
    assets, since a poster overwhelmingly draws on a paper's headline figures.
    """
    by_id = {a['id']: a for a in assets}
    out = [by_id[i] for i in required if i in by_id]
    have = {a['id'] for a in out}
    kind_n: Dict[str, int] = {}
    for a in out:
        kind_n[a['kind']] = kind_n.get(a['kind'], 0) + 1

    for a in sorted(assets, key=lambda x: -x['area']):
        if len(out) >= max_total:
            break
        if a['id'] in have or not a['main_body']:
            continue
        k = a['kind']
        if kind_n.get(k, 0) >= per_kind.get(k, 99):
            continue
        out.append(a)
        have.add(a['id'])
        kind_n[k] = kind_n.get(k, 0) + 1

    # Stable, readable order for the prompt: figures, then tables, then formulas.
    order = {'figure': 0, 'table': 1, 'formula': 2}
    return sorted(out, key=lambda a: (order.get(a['kind'], 9), a['id']))