Blame for app/textconv.py

"""PDF / PMC XML → 标注系统用的「人读 md + layout.json」(1003)。

被两处使用(同一份代码):
  - 01_collection/scripts/make_annotation_text.py   批量转语料
  - 标注服务器的「上传论文」(main.py 以子进程调用本文件的 job 子命令, 依赖 docling/pymupdf 可以只装在转换环境里)

来源优先级: PMC 开放获取子集有 JATS XML → 用 XML(官方原文), 块的 PDF 位置在 PDF 文本层里对齐;
否则 → Docling(版面模型 + TableFormer, 关 OCR), 块位置直接来自 Docling。
layout.json: 每块 md 的 [start, end) + PDF (page, bbox), bbox 为 PDF point、左上原点, page 从 1 起。

子进程用法:
    python textconv.py job --out DIR [--pdf X.pdf] [--ident PMCID|DOI|PMID]
    → DIR/fulltext.md, DIR/layout.json, DIR/paper.pdf, DIR/meta.json(source / pmcid / doi / title / log)
"""
import hashlib
import json
import re
import shutil
import sys
import unicodedata
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET
from pathlib import Path

LAYOUT_VERSION = "1003"


def sha256(p: Path) -> str:
    return hashlib.sha256(p.read_bytes()).hexdigest()


class Builder:
    """顺序拼 md, 同时记块坐标。块之间空一行。"""

    def __init__(self):
        self.parts, self.blocks, self.pos = [], [], 0

    def add(self, md: str, kind: str, boxes=None, **extra):
        md = md.strip("\n")
        if not md.strip():
            return None
        if self.parts:
            self.parts.append("\n\n")
            self.pos += 2
        start = self.pos
        self.parts.append(md)
        self.pos += len(md)
        b = {"id": len(self.blocks), "kind": kind, "start": start, "end": self.pos,
             "boxes": boxes or [], **extra}
        self.blocks.append(b)
        return b

    @property
    def md(self) -> str:
        return "".join(self.parts) + "\n"


def one_line(t: str) -> str:
    return re.sub(r"\s+", " ", t or "").strip()


def cell(t: str) -> str:
    return one_line(t).replace("|", "\\|")


def grid_markdown(grid) -> str:
    """单元格网格 → markdown 表(合并单元格已展开成重复值, 保住列身份)。第一行当表头。"""
    grid = [[cell(c) for c in r] for r in grid if any((c or "").strip() for c in r)]
    if not grid:
        return ""
    w = max(len(r) for r in grid)
    grid = [r + [""] * (w - len(r)) for r in grid]
    lines = ["| " + " | ".join(grid[0]) + " |", "|" + "---|" * w]
    return "\n".join(lines + ["| " + " | ".join(r) + " |" for r in grid[1:]])


# ---------------------------------------------------------------- Docling (PDF only)

CAPTION_START = re.compile(r"(Fig\.?|Figure|FIGURE|Table|TABLE)\s*S?\d+[A-Za-z]?\b")

def convert_docling(pdf: Path):
    from docling.datamodel.base_models import InputFormat
    from docling.datamodel.pipeline_options import PdfPipelineOptions
    from docling.document_converter import DocumentConverter, PdfFormatOption
    from docling_core.types.doc import PictureItem, TableItem

    opts = PdfPipelineOptions()
    opts.do_ocr = False                    # 原生电子版 PDF: 用嵌入文本, 不让 OCR 改字
    opts.do_table_structure = True
    conv = DocumentConverter(format_options={InputFormat.PDF: PdfFormatOption(pipeline_options=opts)})
    doc = conv.convert(str(pdf)).document

    pages = {str(n): [round(p.size.width, 2), round(p.size.height, 2)] for n, p in doc.pages.items()}

    def boxes(item):
        out = []
        for pv in getattr(item, "prov", None) or []:
            h = doc.pages[pv.page_no].size.height
            bb = pv.bbox.to_top_left_origin(page_height=h)
            out.append({"page": pv.page_no, "bbox": [round(bb.l, 1), round(bb.t, 1), round(bb.r, 1), round(bb.b, 1)]})
        return out

    B = Builder()
    # 归属表/图的 caption 跟着表/图输出; 正文顺序里单独遇到时跳过, 否则会出现两遍
    owned = {c.cref for it in list(doc.tables) + list(doc.pictures) for c in it.captions}
    for item, _lvl in doc.iterate_items(traverse_pictures=False):
        if item.self_ref in owned:
            continue
        label = item.label.value
        if isinstance(item, (TableItem, PictureItem)):
            caps = [cref.resolve(doc) for cref in item.captions]
            # Docling 有时把图注挂在表 / 图的子节点里、不登记进 captions(traverse_pictures=False 会整个跳过), 这里补回(1016)。
            # 只收 label=caption, 或以 "Fig./Figure/Table + 编号" 开头的子节点文字(会被错标成 footnote / text);
            # 图内其余文字(坐标轴、图例)不收。
            seen = {c.self_ref for c in caps}
            for ch, _ in doc.iterate_items(root=item, traverse_pictures=True):
                txt = one_line(getattr(ch, "text", ""))
                if ch is item or ch.self_ref in seen or not txt:
                    continue
                if ch.label.value == "caption" or CAPTION_START.match(txt):
                    caps.append(ch)
                    seen.add(ch.self_ref)
                    owned.add(ch.self_ref)         # 主循环之后若再遇到同一节点, 不再输出第二遍
            for cap in caps:
                B.add(one_line(cap.text), "caption", boxes(cap))
            if isinstance(item, TableItem):
                md = grid_markdown([[c.text for c in row] for row in item.data.grid])
                B.add(md or "[table — not recovered, see PDF]", "table", boxes(item))
            else:
                B.add("[figure — see PDF]", "figure", boxes(item))
            continue
        text = one_line(getattr(item, "text", ""))
        if label == "text" and re.fullmatch(r"\d{1,4}", text):
            continue                       # 投稿版页边行号(000 001 …), 每个都会被识别成一块
        link = getattr(item, "hyperlink", None)
        if link and not str(link).startswith("mailto:") and str(link) not in text:
            text = f"{text} <{link}>"      # PyMuPDF 版只剩锚词的超链接, 这里把目标补回来
        if label == "title":
            B.add("# " + text, "heading", boxes(item))
        elif label == "section_header":
            lv = min(max(getattr(item, "level", 1), 1), 4)
            B.add("#" * (lv + 1) + " " + text, "heading", boxes(item))
        elif label == "list_item":
            mk = (getattr(item, "marker", "") or "").strip()     # 保留原编号 "(2)" / "a." ; 圆点类统一成 "-"
            B.add(("- " if not mk or mk in "-•·▪◦–*" else mk + " ") + text, "list", boxes(item))
        elif label == "formula":
            B.add(text or "[formula — not recovered, see PDF]", "formula", boxes(item))
        elif label == "code":
            B.add("```\n" + (getattr(item, "text", "") or "") + "\n```", "code", boxes(item))
        elif label in ("footnote", "caption", "reference"):
            B.add(text, label, boxes(item))
        else:
            B.add(text, "text", boxes(item))
    return B, pages, {"source": "docling", "docling_version": _ver("docling")}


def _ver(pkg):
    try:
        import importlib.metadata as m
        return m.version(pkg)
    except Exception:
        return None


# ---------------------------------------------------------------- JATS XML (官方原文)

def tag(el):
    return el.tag.split("}", 1)[-1]


XLINK = "{http://www.w3.org/1999/xlink}href"
NESTED_BLOCKS = {"table-wrap", "fig", "fig-group", "table-wrap-group", "disp-formula", "list", "boxed-text"}


def child(el, name):
    return el.find("{*}" + name)                  # {*} 同时匹配有 / 无命名空间


def inline(el) -> str:
    """段落内联文本: 保留 ext-link 的 URL、行内公式(tex 或 MathML 文本)、上下标(^ _)。段内嵌的表/图/列表由调用方单独输出。"""
    out = [el.text or ""]
    for ch in el:
        t = tag(ch)
        if t in NESTED_BLOCKS:
            pass
        elif t in ("ext-link", "uri"):
            txt = one_line("".join(ch.itertext()))
            href = ch.get(XLINK) or ""
            out.append(txt + (f" <{href}>" if href and href not in txt else ""))
        elif t == "inline-formula":
            tex = ch.find(".//{*}tex-math")
            out.append(one_line(tex.text if tex is not None and tex.text else "".join(ch.itertext())))
        elif t == "sup":
            out.append("^" + "".join(ch.itertext()))
        elif t == "sub":
            out.append("_" + "".join(ch.itertext()))
        else:
            out.append(inline(ch))
        out.append(ch.tail or "")
    return "".join(out)


def jats_table(tw) -> str:
    """table-wrap → markdown 表; 处理 colspan / rowspan(展开成重复单元格, 保住列身份)。"""
    rows = [tr for tr in tw.iter() if tag(tr) == "tr"]
    grid, carry = [], {}
    for tr in rows:
        row, col = [], 0
        cells = [c for c in tr if tag(c) in ("td", "th")]
        ci = 0
        while ci < len(cells) or col in carry:
            if col in carry:
                txt, left = carry[col]
                row.append(txt)
                if left <= 1:
                    del carry[col]
                else:
                    carry[col] = (txt, left - 1)
                col += 1
                continue
            c = cells[ci]
            ci += 1
            txt = one_line(inline(c))          # grid_markdown 统一转义
            cs, rs = int(c.get("colspan", 1) or 1), int(c.get("rowspan", 1) or 1)
            for k in range(cs):
                row.append(txt)
                if rs > 1:
                    carry[col + k] = (txt, rs - 1)
            col += cs
        grid.append(row)
    return grid_markdown(grid)


def convert_jats(xml: Path):
    root = ET.parse(xml).getroot()
    B = Builder()
    # 很多 PMC XML 把表/图集中放在 <floats-group>: 插到正文第一次引用它的段落后面, 没被引用的放正文末尾
    fg = next((x for x in root if tag(x) == "floats-group"), None)
    floats = {f.get("id"): f for f in (fg if fg is not None else []) if f.get("id")}
    placed = set()

    def place_floats(p):
        for x in p.iter():
            if tag(x) == "xref" and x.get("ref-type") in ("table", "fig"):
                for rid in (x.get("rid") or "").split():
                    if rid in floats and rid not in placed:
                        placed.add(rid)
                        blocks_of(floats[rid], 1)

    def caption_of(el):
        lab, cap = child(el, "label"), child(el, "caption")
        parts = [one_line(inline(lab)) if lab is not None else ""]
        if cap is not None:
            parts += [one_line(inline(x)) for x in cap if tag(x) in ("title", "p")]
        return one_line(" ".join(p for p in parts if p))

    def blocks_of(el, depth):
        t = tag(el)
        if t == "sec" or t == "app" or t == "ack":
            title, label = child(el, "title"), child(el, "label")
            head = one_line(" ".join(x for x in [inline(label) if label is not None else "",
                                                  inline(title) if title is not None else ""] if x))
            if t == "ack" and not head:
                head = "Acknowledgments"
            if head:
                B.add("#" * min(depth + 1, 5) + " " + head, "heading")
            for ch in el:
                if tag(ch) not in ("title", "label"):
                    blocks_of(ch, depth + 1)
        elif t == "p":
            B.add(one_line(inline(el)), "text")
            for ch in el.iter():
                if ch is not el and tag(ch) in ("table-wrap", "fig", "disp-formula", "list"):
                    blocks_of(ch, depth)
            place_floats(el)
        elif t == "list":
            for li in el:
                if tag(li) == "list-item":
                    B.add("- " + one_line(" ".join(inline(p) for p in li if tag(p) == "p")), "list")
        elif t == "table-wrap":
            B.add(caption_of(el), "caption")
            for tb in el.iter():
                if tag(tb) == "table":
                    B.add(jats_table(tb) or "[table — not in XML]", "table")
            for fn in el.iter():
                if tag(fn) == "table-wrap-foot":
                    B.add(one_line(inline(fn)), "footnote")
        elif t == "fig":
            B.add(caption_of(el), "caption")
            B.add("[figure — see PDF]", "figure")
        elif t == "disp-formula":
            tex = el.find(".//{*}tex-math")
            B.add(one_line(tex.text if tex is not None and tex.text else "".join(el.itertext())) or
                  "[formula — see PDF]", "formula")
        elif t in ("table-wrap-group", "fig-group", "boxed-text"):
            for ch in el:
                blocks_of(ch, depth)
        elif t == "fn-group":
            for fn in el:
                B.add(one_line(" ".join(inline(p) for p in fn.iter() if tag(p) == "p")), "footnote")
        elif t == "supplementary-material":
            href = el.get(XLINK) or ""
            B.add(one_line(caption_of(el) + (f" <{href}>" if href else "")) or "[supplementary material]",
                  "supplementary")
        elif t == "ref-list":
            B.add("## References", "heading")
            for r in el.iter():
                if tag(r) == "ref":
                    c = next((x for x in r if tag(x) in ("mixed-citation", "element-citation", "citation")), None)
                    lab = child(r, "label")
                    txt = one_line(" ".join(c.itertext())) if c is not None else ""
                    B.add(one_line(((lab.text or "") + " " if lab is not None else "") + txt), "reference")
        elif t in ("title", "label"):
            pass
        else:
            for ch in el:
                blocks_of(ch, depth)

    front = next((x for x in root if tag(x) == "front"), None)
    if front is not None:
        title = next((x for x in front.iter() if tag(x) == "article-title"), None)
        if title is not None:
            B.add("# " + one_line(inline(title)), "heading")
        for ab in front.iter():
            if tag(ab) == "abstract":
                B.add("## Abstract" + (f" ({ab.get('abstract-type')})" if ab.get("abstract-type") else ""), "heading")
                for ch in ab:
                    blocks_of(ch, 2)
    body = next((x for x in root if tag(x) == "body"), None)
    if body is not None:
        for ch in body:
            blocks_of(ch, 1)
    for rid, f in floats.items():
        if rid not in placed:
            placed.add(rid)
            blocks_of(f, 1)
    back = next((x for x in root if tag(x) == "back"), None)
    if back is not None:
        for ch in back:
            blocks_of(ch, 1)
    return B, {"source": "jats", "xml_sha256": sha256(xml)}


# ---------------------------------------------------------------- 对齐: md 块 → PDF 文本层

def norm_chars(s: str):
    """NFKC + 小写 + 只留字母数字: 对断词连字符、空格、连字(fi)都不敏感。"""
    s = unicodedata.normalize("NFKC", s).lower()
    return "".join(ch for ch in s if ch.isalnum())


def align_to_pdf(B: Builder, md: str, pdf: Path):
    import pymupdf as fitz
    doc = fitz.open(pdf)
    pages = {}
    stream, owner = [], []                 # 全文归一化字符流; 每个字符属于哪个 word
    words = []
    for pno, page in enumerate(doc, start=1):
        pages[str(pno)] = [round(page.rect.width, 2), round(page.rect.height, 2)]
        for w in page.get_text("words"):
            x0, y0, x1, y1, txt, blk, ln, _ = w
            n = norm_chars(txt)
            if not n:
                continue
            wi = len(words)
            words.append((pno, blk, (x0, y0, x1, y1)))
            stream.append(n)
            owner.extend([wi] * len(n))
    doc.close()
    S = "".join(stream)
    cursor, hit = 0, 0
    for b in B.blocks:
        if b["kind"] == "figure":
            continue
        s = norm_chars(re.sub(r"<[^>]+>|\|---|[#|*]", " ", md[b["start"]:b["end"]]))
        if len(s) < 8:
            continue
        probe = s[:40]
        i = S.find(probe, cursor)
        if i < 0:
            i = S.find(probe)
        if i < 0 and len(s) > 80:            # 开头对不上(如表头次序不同)就试中段
            mid = len(s) // 2
            j = S.find(s[mid:mid + 40])
            i = j - mid if j >= 0 else -1
        if i < 0:
            continue
        i = max(i, 0)
        tail = s[-40:]
        j = S.find(tail, i, i + int(len(s) * 1.5) + 400)
        end = j + len(tail) if j >= 0 else min(i + len(s), len(S))
        cursor = i + 1
        groups = {}
        for wi in sorted(set(owner[i:end])):
            pno, blk, (x0, y0, x1, y1) = words[wi]
            g = groups.setdefault((pno, blk), [x0, y0, x1, y1])
            g[0], g[1], g[2], g[3] = min(g[0], x0), min(g[1], y0), max(g[2], x1), max(g[3], y1)
        b["boxes"] = [{"page": p, "bbox": [round(v, 1) for v in bb]} for (p, _), bb in groups.items()]
        hit += 1
    return pages, hit


# ---------------------------------------------------------------- 输出

def layout_dict(B: Builder, pages: dict, meta: dict, pdf) -> dict:
    return {"layout_version": LAYOUT_VERSION, **meta, "pdf_sha256": sha256(Path(pdf)) if pdf else None,
            "md_sha256": hashlib.sha256(B.md.encode("utf-8")).hexdigest(), "pages": pages, "blocks": B.blocks}


def md_title(B: Builder):
    for b in B.blocks:
        if b["kind"] == "heading":
            t = B.md[b["start"]:b["end"]].lstrip("#").strip()
            if len(t) > 12 and not re.match(r"^(abstract|introduction|\d+\s)", t, re.I):
                return t[:300]
    return None


# ---------------------------------------------------------------- PMC: 识别 + 拉取(开放获取子集)

UA = "reproducibility-annotator/1.0 (research; mailto:none)"
IDCONV = "https://pmc.ncbi.nlm.nih.gov/tools/idconv/api/v1/articles/?format=json&tool=repro-annotator&ids="
OA = "https://pmc-oa-opendata.s3.amazonaws.com"
DOI_RE = re.compile(r"\b(10\.\d{4,9}/[^\s\"<>]+)", re.I)


def _get(url: str, timeout=60) -> bytes:
    req = urllib.request.Request(url, headers={"User-Agent": UA})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        return r.read()


def to_pmcid(ident: str, log):
    """PMCID / DOI / PMID → PMCID(查不到返回 None)。"""
    s = ident.strip()
    m = re.fullmatch(r"(?i)(pmc)?(\d{5,9})", s)
    if m and m.group(1):
        return "PMC" + m.group(2)
    s = re.sub(r"(?i)^(https?://(dx\.)?doi\.org/|doi:\s*)", "", s)
    try:
        recs = json.loads(_get(IDCONV + urllib.parse.quote(s, safe="")))["records"]
    except Exception as e:  # noqa: BLE001
        log(f"PMC ID lookup failed for {s}: {e}")
        return None
    pmcid = next((r.get("pmcid") for r in recs if r.get("pmcid")), None)
    log(f"{s} → {pmcid or 'not in PMC'}")
    return pmcid


def ident_from_pdf(pdf: Path, log):
    """PDF 前两页里找 PMCID(作者稿常印在页眉)或 DOI。"""
    import pymupdf as fitz
    with fitz.open(pdf) as d:
        text = " ".join(d[i].get_text() for i in range(min(2, len(d))))
    m = re.search(r"\bPMC\d{5,9}\b", text)
    if m:
        log(f"found {m.group(0)} printed in the PDF")
        return m.group(0)
    m = DOI_RE.search(text)
    if m:
        doi = m.group(1).rstrip(".,;)]")
        log(f"found DOI {doi} in the PDF")
        return doi
    log("no PMCID / DOI found in the PDF")
    return None


def pmc_oa(pmcid: str, log):
    """开放获取子集里这篇的最新版本元数据(含 xml_url / pdf_url / license); 不在 OA 子集返回 None。"""
    listing = _get(f"{OA}/?list-type=2&delimiter=/&prefix={pmcid}.").decode()
    vers = sorted(int(v) for v in re.findall(rf"<Prefix>{pmcid}\.(\d+)/</Prefix>", listing))
    if not vers:
        log(f"{pmcid} is not in the PMC open-access subset (no XML available)")
        return None
    key = f"{pmcid}.{vers[-1]}"
    meta = json.loads(_get(f"{OA}/{key}/{key}.json"))
    fix = lambda u: (u or "").split("?")[0].replace("s3://pmc-oa-opendata", OA) or None  # noqa: E731
    meta["xml_url"], meta["pdf_url"] = fix(meta.get("xml_url")), fix(meta.get("pdf_url"))
    log(f"{key}: open access, license {meta.get('license_code')}, xml={'yes' if meta['xml_url'] else 'no'}, "
        f"pdf={'yes' if meta['pdf_url'] else 'no'}")
    return meta


def run_job(out: Path, pdf=None, ident=None) -> dict:
    out.mkdir(parents=True, exist_ok=True)
    logs = []
    log = lambda m: (logs.append(m), print(m, flush=True))  # noqa: E731
    info = {"source": None, "pmcid": None, "doi": None, "title": None, "license": None}
    pdf = Path(pdf) if pdf else None

    guess = ident or (ident_from_pdf(pdf, log) if pdf else None)
    if guess and re.match(r"10\.", re.sub(r"(?i)^(https?://(dx\.)?doi\.org/|doi:\s*)", "", guess.strip())):
        info["doi"] = re.sub(r"(?i)^(https?://(dx\.)?doi\.org/|doi:\s*)", "", guess.strip())
    pmcid = to_pmcid(guess, log) if guess else None
    oa = pmc_oa(pmcid, log) if pmcid else None
    info["pmcid"] = pmcid

    if oa and oa.get("xml_url"):
        xml = out / "source.xml"
        xml.write_bytes(_get(oa["xml_url"]))
        if not pdf:
            if not oa.get("pdf_url"):
                raise SystemExit("PMC has XML but no PDF for this article; please upload the PDF")
            (out / "paper.pdf").write_bytes(_get(oa["pdf_url"], timeout=180))
            log("downloaded the PDF from PMC")
        else:
            shutil.copyfile(pdf, out / "paper.pdf")
        B, meta = convert_jats(xml)
        pages, hit = align_to_pdf(B, B.md, out / "paper.pdf")
        log(f"PMC XML → {len(B.blocks)} blocks, {hit} located on the PDF")
        info.update(source="jats", title=oa.get("title") or md_title(B), license=oa.get("license_code"))
        meta["pmcid"] = pmcid
    else:
        if not pdf:
            raise SystemExit("no PDF uploaded and no open-access PMC XML found")
        shutil.copyfile(pdf, out / "paper.pdf")
        log("converting the PDF with Docling (layout + tables)…")
        B, pages, meta = convert_docling(out / "paper.pdf")
        log(f"Docling → {len(pages)} pages, {len(B.blocks)} blocks")
        info.update(source="docling", title=md_title(B))

    (out / "fulltext.md").write_text(B.md, encoding="utf-8")
    (out / "layout.json").write_text(json.dumps(layout_dict(B, pages, meta, out / "paper.pdf"), ensure_ascii=False),
                                     encoding="utf-8")
    info["log"] = logs
    (out / "meta.json").write_text(json.dumps(info, ensure_ascii=False, indent=1), encoding="utf-8")
    return info


if __name__ == "__main__":
    import argparse
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd", choices=["job"])
    ap.add_argument("--out", required=True)
    ap.add_argument("--pdf")
    ap.add_argument("--ident")
    a = ap.parse_args()
    try:
        run_job(Path(a.out), a.pdf, a.ident)
    except SystemExit as e:
        print(f"ERROR: {e}", flush=True)
        sys.exit(2)