SPB Git

Repository Public Research preview

pairs.py 3.7 KB 64 lines python

1"""Build validated pairs (source = rewrite, target = original) from a rewrite job and the corpus."""2from __future__ import annotations34from collections import defaultdict5from pathlib import Path67from htr.data.schema import Pair8from htr.data.taxonomy import era_of9from htr.parse.protect import protect10from htr.utils.hashing import short_id11from htr.utils.io import read_jsonl, write_json, write_jsonl12from htr.utils.provenance import run_header13from htr.validate.checks import VALIDATOR_VERSION, ValidationConfig, validate_pair141516def build_pairs(corpus_dir: Path, job_dir: Path, out_dir: Path, cfg: ValidationConfig | None = None) -> dict:17    cfg = cfg or ValidationConfig()18    docs = {d["doc_id"]: d for d in read_jsonl(corpus_dir / "documents.jsonl")}19    passages = {p["passage_id"]: p for p in read_jsonl(corpus_dir / "passages.jsonl")}20    pairs: list[Pair] = []21    stats: dict = {"rewrites": 0, "status": defaultdict(int), "validated": 0, "passed": 0, "quarantined": 0, "by_variant": defaultdict(lambda: {"n": 0, "passed": 0}), "fail_reasons": defaultdict(int)}22    for rw in read_jsonl(job_dir / "rewrites.jsonl"):23        stats["rewrites"] += 124        stats["status"][rw["status"]] += 125        p = passages.get(rw["passage_id"])26        if p is None:27            continue28        d = docs[p["doc_id"]]29        prot = protect(p["text"], version=rw.get("protect_version", "1.0"))  # same placeholder protocol as the request30        val = validate_pair(p["text"], rw.get("text"), rw["variant"], cfg, original_protected=prot.text) if rw["status"] == "ok" else None31        passed = bool(val and val.passed)32        stats["validated"] += 133        stats["by_variant"][rw["variant"]]["n"] += 134        if passed:35            stats["passed"] += 136            stats["by_variant"][rw["variant"]]["passed"] += 137        else:38            stats["quarantined"] += 139            for r in (val.reasons if val else [f"status:{rw['status']}"]):40                stats["fail_reasons"][r.split(" ")[0] + (" " + r.split(" ")[1] if r.startswith(("length", "numbers", "placeholders", "entity", "semantic", "sentence")) else "")] += 141        pairs.append(42            Pair(43                pair_id=short_id("pair", rw["rewrite_id"]),44                passage_id=p["passage_id"],45                doc_id=p["doc_id"],46                rewrite_id=rw["rewrite_id"],47                variant=rw["variant"],48                sample_index=rw["sample_index"],49                source_text=rw.get("text") or "",50                target_text=prot.text,  # placeholder form; original spans restorable via passage.protected51                metadata={"discipline": d.get("subject_area"), "pub_year": d.get("pub_year"), "era": era_of(d.get("pub_year")), "section_role": p["section_role"], "source": d["source"], "license": d["license"]["code"], "venue": d.get("venue"), "author_group": d.get("author_group"), "model": rw["model"], "prompt_version": rw["prompt_version"], "protect_version": rw.get("protect_version", "1.0"), "n_words_target": p["n_words"]},52                validation=val,53                quarantined=not passed,54            )55        )56    out_dir.mkdir(parents=True, exist_ok=True)57    write_jsonl(out_dir / "pairs.jsonl", [p for p in pairs if not p.quarantined])58    write_jsonl(out_dir / "quarantine.jsonl", [p for p in pairs if p.quarantined])59    stats = {k: (dict(v) if isinstance(v, defaultdict) else v) for k, v in stats.items()}60    stats["by_variant"] = {k: dict(v) for k, v in stats["by_variant"].items()}61    stats["acceptance_rate"] = round(stats["passed"] / max(1, stats["validated"]), 4)62    write_json(out_dir / "validation_summary.json", {**run_header({"validator_version": VALIDATOR_VERSION, "config": cfg.__dict__}), "stats": stats})63    return stats64
Files main
Menu

Theme