PlantAI-rag-new/tooling/aqua_kb/qc_sample.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""qc_sample.py — 抽取 manifest.csv 三种状态样本并做可核验的规则检查。

用法: qc_sample.py [--manifest build/manifest.csv] [--source aqua/source]
                    [--n 20] [--seed 42]

规则(任一不通过 exit 1):
  paired: md_path、pdf_path 都存在;md 文本 grep 到 doi_norm(斜杠或下划线形态
          任一命中);PDF 文件名归一化后等于该 doi_norm。
  md_only: pdf_path 为空;PDF正文 目录下无 *<doi_norm 下划线形态>* 与
           *<斜杠换成下划线>* 文件;DOI汇总.txt 中 grep 到该 DOI(两种形态)。
  pdf_only: md_path 为空;题目摘要_MD 目录下 grep 不到该 DOI(两种形态)。

只依赖标准库。
"""
from __future__ import annotations

import argparse
import csv
import random
import re
import sys
from pathlib import Path


def norm_doi(d: str) -> str:
    """DOI 归一化(与 pair_aqua_sources 同规则)。"""
    return d.strip().lower().rstrip(".,;)]")


def slug_form(doi_norm: str) -> str:
    """斜杠换下划线 + 点换下划线:文件名 DOI 段的下划线形态。"""
    return doi_norm.replace("/", "_").replace(".", "_")


def grep_pdf_norm(filename: str, doi_norm: str) -> bool:
    """PDF 文件名归一化后是否等于 doi_norm(复用 pair 的解码算法,取等值判定)。"""
    stem = filename[:-4] if filename.lower().endswith(".pdf") else filename
    cands: list[str] = []
    for pos in reversed([m.start() for m in re.finditer(r"10\.", stem)]):
        segs = stem[pos + 3:].split("_")
        if not segs or not re.fullmatch(r"\d{4,9}", segs[0]):
            continue
        parts = [segs[0]]
        for i in range(1, len(segs)):
            seg = segs[i]
            if i == len(segs) - 1 and re.fullmatch(r"[0-9a-fA-F]{8}", seg):
                break
            if seg and seg[0].isalnum():
                parts.append(seg)
            else:
                break
        for k in range(1, len(parts) + 1):
            cands.append(norm_doi("10." + "/".join(parts[:k])))
    return doi_norm in cands


def grep_file_text(f: Path, needle_slash: str, needle_under: str) -> bool:
    if not f.is_file():
        return False
    text = f.read_text(encoding="utf-8", errors="replace")
    return needle_slash in text or needle_under in text


def main() -> int:
    ap = argparse.ArgumentParser(description="manifest 状态样本质量抽检")
    ap.add_argument("--manifest", default="build/manifest.csv")
    ap.add_argument("--source", default="aqua/source")
    ap.add_argument("--n", type=int, default=20)
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    manifest = Path(args.manifest)
    src = Path(args.source)
    md_dir = src / "题目摘要_MD"
    pdf_dir = src / "PDF正文"
    doi_txt = src / "DOI汇总.txt"

    rows = list(csv.DictReader(manifest.open(encoding="utf-8", newline="")))
    groups: dict[str, list[dict]] = {"paired": [], "md_only": [], "pdf_only": []}
    for r in rows:
        if r["status"] in groups:
            groups[r["status"]].append(r)

    rng = random.Random(args.seed)
    samples = {st: rng.sample(groups[st], min(args.n, len(groups[st]))) for st in groups}

    def check_md_text(doi_norm: str, md_path: str) -> bool:
        text = (src / md_path).read_text(encoding="utf-8", errors="replace")
        return doi_norm in text or slug_form(doi_norm) in text

    def check_md_dir_grep(doi_norm: str) -> bool:
        n_slash = doi_norm
        n_under = slug_form(doi_norm)
        for f in md_dir.glob("*.md"):
            t = f.read_text(encoding="utf-8", errors="replace")
            if n_slash in t or n_under in t:
                return True
        return False

    fails = 0
    checked: dict[str, int] = {}

    print("== qc_sample ==")
    print(f"manifest={manifest} rows={len(rows)} "
          f"(paired={len(groups['paired'])} md_only={len(groups['md_only'])} "
          f"pdf_only={len(groups['pdf_only'])}) seed={args.seed} n={args.n}\n")

    for st in ("paired", "md_only", "pdf_only"):
        print(f"--- {st} (抽 {len(samples[st])} 条) ---")
        for r in samples[st]:
            dn = r["doi_norm"]; md = r["md_path"]; pdf = r["pdf_path"]
            ok = True
            if st == "paired":
                if not (md and pdf and (src / md).is_file() and (src / pdf).is_file()):
                    ok = False
                elif not check_md_text(dn, md):
                    ok = False
                elif not grep_pdf_norm(Path(pdf).name, dn):
                    ok = False
            elif st == "md_only":
                if pdf:
                    ok = False
                else:
                    has_glob = any(
                        pdf_dir.glob("*" + slug_form(dn) + "*")
                    ) or any(
                        pdf_dir.glob("*" + dn.replace("/", "_") + "*")
                    )
                    if has_glob:
                        ok = False
                    elif not grep_file_text(doi_txt, dn, slug_form(dn)):
                        ok = False
            else:  # pdf_only
                if md:
                    ok = False
                elif check_md_dir_grep(dn):
                    ok = False
            if not ok:
                fails += 1
            print(f"[{'PASS' if ok else 'FAIL'}] {st:8s} "
                  f"doi={dn[:46]:46s} md={Path(md).name if md else '-':20s} "
                  f"pdf={Path(pdf).name if pdf else '-'}")
        checked[st] = len(samples[st])

    print("\n== 结果 ==")
    print(f"checked={sum(checked.values())} fails={fails}")
    for st in ("paired", "md_only", "pdf_only"):
        print(f"  {st}: {checked[st]} 条")
    print("PASS" if fails == 0 else "FAIL")
    return 1 if fails else 0


if __name__ == "__main__":
    sys.exit(main())