bid-doctor/spikes/s4-llm/prep.py

"""S4 数据准备：把规范化 PDF 切成带编号、带格式标记的句子，并生成银标准。

句子编号和 S3 一样（同一套切句规则），格式标记：
  斜体且下划线的连续文字 → <iu>…</iu>；只有下划线 → <u>…</u>；只有斜体 → <i>…</i>
银标准（实质性条款句）：
  ★ 句：句子含 ★（排除「带 ★ 的为实质性要求」这类说明句，以及「★号条款」「“★”项」这类引用句）
  斜体下划线句：句子里一半以上的字是斜体且下划线，且文件自己声明了这种标注方式（S06 / S08 / S09）
另外抽出「标注规则说明」句（文件自己说明怎么标注实质性条款的句子），作为每个分块的背景。

用法：python -I prep.py <规范化 PDF 目录> S01 S02 …
"""
from __future__ import annotations

import json
import pathlib
import re
import sys

import pdfplumber

HERE = pathlib.Path(__file__).resolve().parent
sys.path.insert(0, str(HERE.parent / "s3-citation"))
from sentences import split_sentences  # noqa: E402  同一套切句规则，保证编号一致

OUT = HERE / "out"
WS = re.compile(r"[\s　]")
# 说明标注规则的句子：既提到标注方式，又提到实质性 / 无效
# 只是「提到」★ 条款的句子（“★”、★号、★项、★条款），本身不是条款
STAR_REF = re.compile(r"[“\"]★[”\"]|★\s*(号|项|条款|标注|标记)")
RULE_NOTE = re.compile(r"(★|▲|\*|斜体|下划线|加粗|标注|标记|加注|标有).{0,30}(实质性|无效|否决|废标)|(实质性|无效|否决).{0,30}(★|斜体|下划线|标注|标记|加注|标有)")


def char_flags(page):
    hlines = [(l["x0"], l["x1"], l["top"]) for l in page.lines if abs(l["top"] - l["bottom"]) < 0.5]
    hlines += [(r["x0"], r["x1"], r["top"]) for r in page.rects if r["height"] < 1.5 and r["width"] > 2]
    flags = {}
    for c in page.chars:
        if not c["text"].strip():
            continue
        name = c["fontname"].split("+")[-1]
        it = "Italic" in name or "Oblique" in name or abs(c["matrix"][2]) > 0.1
        mid = (c["x0"] + c["x1"]) / 2
        un = any(x0 - 0.5 <= mid <= x1 + 0.5 and -0.25 * c["size"] <= y - c["bottom"] <= 4 for x0, x1, y in hlines)
        flags[(round(c["x0"], 2), round(c["top"], 2), c["text"])] = (it, un)
    return flags


def chars_in_quads(page, quads):
    got = []
    for c in page.chars:
        if not c["text"].strip():
            continue
        cx, cy = (c["x0"] + c["x1"]) / 2, (c["top"] + c["bottom"]) / 2
        for qi, (x0, t, x1, b) in enumerate(quads):
            if x0 <= cx <= x1 and t <= cy <= b:
                got.append((qi, c["x0"], c)); break
    got.sort(key=lambda x: (x[0], x[1]))
    return [c for _, _, c in got]


def markup(chars, flags):
    out, cur_tag = [], None
    for c in chars:
        it, un = flags.get((round(c["x0"], 2), round(c["top"], 2), c["text"]), (False, False))
        tag = "iu" if it and un else "u" if un else "i" if it else None
        if tag != cur_tag:
            if cur_tag:
                out.append(f"</{cur_tag}>")
            if tag:
                out.append(f"<{tag}>")
            cur_tag = tag
        out.append(c["text"])
    if cur_tag:
        out.append(f"</{cur_tag}>")
    iu = sum(1 for c in chars if flags.get((round(c["x0"], 2), round(c["top"], 2), c["text"]), (0, 0)) == (True, True))
    return "".join(out), iu


def main():
    pdf_dir, sids = pathlib.Path(sys.argv[1]), sys.argv[2:]
    for sid in sids:
        sents = []
        with pdfplumber.open(pdf_dir / f"{sid}.pdf") as pdf:
            for page in pdf.pages:
                right = max((c["x1"] for c in page.chars), default=page.width)
                flags = char_flags(page)
                for s in split_sentences(sid, page, right):
                    chars = chars_in_quads(page, s["quads"])
                    marked, iu = markup(chars, flags)
                    plain = "".join(c["text"] for c in chars)
                    n = len(WS.sub("", plain))
                    is_rule_note = bool(RULE_NOTE.search(plain))
                    silver = None
                    if "★" in plain and not is_rule_note and not STAR_REF.search(plain):
                        silver = "star"
                    elif n >= 6 and iu / max(1, n) >= 0.5 and not is_rule_note:
                        silver = "iu"
                    sents.append({"id": s["id"].split(":", 1)[1].replace(":", ""),  # p12s3
                                  "page": s["page"], "text": plain, "marked": marked,
                                  "silver": silver, "rule_note": is_rule_note})
        # 斜体下划线只在文件自己声明了这种标注方式时才算实质性（其他文件里它常是模板填空处）
        declares_iu = any(s["rule_note"] and "斜体" in s["text"] for s in sents)
        for s in sents:
            if s["silver"] == "iu" and not declares_iu:
                s["silver"] = None
        (OUT / f"{sid}-input.json").write_text(json.dumps(sents, ensure_ascii=False))
        print(sid, "sentences", len(sents), "silver_star", sum(s["silver"] == "star" for s in sents),
              "silver_iu", sum(s["silver"] == "iu" for s in sents), "rule_notes", sum(s["rule_note"] for s in sents))


if __name__ == "__main__":
    main()