rafmacalaba/data-use-annotate

1#!/usr/bin/env python32"""Build an annotation queue from rafmacalaba/datause-ner via datasets-server.3 4Stdlib only (urllib) so it runs anywhere, including Spaces builds.5 6    uv run python human_labeling/build_pool_queue.py [--config probe_candidates] [--split pool]7        [--limit 500] [--out human_labeling/queue.json]8 9Defaults target the annotatable config: probe_candidates/pool is the only10one with an embedded `passage`. --config probe_splits yields surface-only11triage items flagged ctx_missing (no passage on the Hub; see README).12Stratified sample: round-robin over origin x score-band so the queue13covers keep/confusion/drop and every origin instead of the head-heavy14random draw (proxy bands on the Hub score; true tags come from rescore).15"""16 17import argparse18import json19import urllib.parse20import urllib.request21from pathlib import Path22 23REPO = Path(__file__).resolve().parent24DATASET = "rafmacalaba/datause-ner"25API = "https://datasets-server.huggingface.co/rows"26 27BANDS = (("drop", 0.0, 0.05), ("confusion", 0.05, 0.9), ("keep", 0.9, 1.01))28# Proxy strata on the Hub (v3-era) score, used ONLY to stratify the sample29# across the eventual decision zones (same 0.05/0.9 edges as the audit30# deck). True keep/confusion/drop labels are assigned later by31# rescore_singlepass.probe_labels.decide; the written item band is always32# "unscored" until then.33def band_of(score: float) -> str:34    for name, lo, hi in BANDS:35        if lo <= score < hi:36            return name37    return "confusion"38 39 40def fetch_rows(config: str, split: str, offset: int, length: int) -> list[dict]:41    import time42    import urllib.error43    q = urllib.parse.urlencode(44        {"dataset": DATASET, "config": config, "split": split,45         "offset": offset, "length": length}46    )47    last = None48    for attempt in range(6):49        try:50            with urllib.request.urlopen(f"{API}?{q}", timeout=60) as r:51                payload = json.load(r)52            return [d["row"] for d in payload.get("rows", [])]53        except urllib.error.HTTPError as e:54            last = e55            if e.code not in (429, 500, 502, 503):56                raise57            time.sleep(2 ** attempt)58    raise last59def to_item(row: dict, config: str, split: str) -> dict:60    score = float(row.get("head_score") or 0.0)61    passage = row.get("passage")62    start, end = row.get("start"), row.get("end")63    item = {64        "key": row.get("key"),65        "surface": row.get("surface"),66        "ctx": passage,  # None for probe_splits rows (no passage on Hub)67        "ctx_missing": passage is None,68        "head_score": score,  # Hub v3-era placeholder; rescore overwrites69        "start": start, "end": end,70        "band": "unscored",  # rescore_singlepass tags keep/confusion/drop71        "origin": row.get("origin"),72        "specificity": row.get("extractor_label") or row.get("stratum") or "",73        "split": row.get("split", split),74        "queue": f"{config}/{split}",75    }76    return item77 78 79def main() -> None:80    ap = argparse.ArgumentParser()81    ap.add_argument("--config", default="probe_candidates")82    ap.add_argument("--split", default="pool")83    ap.add_argument("--limit", type=int, default=500)84    ap.add_argument("--fetch", type=int, default=5000,85                    help="rows scanned from the Hub for stratification")86    ap.add_argument("--seed", type=int, default=7)87    ap.add_argument("--out", default=str(REPO / "queue.json"))88    a = ap.parse_args()89 90    scanned: list[dict] = []91    offset = 092    while len(scanned) < a.fetch:93        batch = fetch_rows(a.config, a.split, offset, min(100, a.fetch - len(scanned)))94        if not batch:95            break96        scanned.extend(batch)97        offset += len(batch)98 99    # stratify: round-robin over (origin, band)100    import random101    rng = random.Random(a.seed)102    buckets: dict[tuple, list[dict]] = {}103    for row in scanned:104        item = to_item(row, a.config, a.split)105        buckets.setdefault((item["origin"], item["band"]), []).append(item)106    for b in buckets.values():107        rng.shuffle(b)108    queue: list[dict] = []109    while len(queue) < min(a.limit, len(scanned)) and buckets:110        for k in sorted(buckets):111            if len(queue) >= a.limit:112                break113            if buckets[k]:114                queue.append(buckets[k].pop())115        buckets = {k: v for k, v in buckets.items() if v}116    rng.shuffle(queue)117 118    out = Path(a.out)119    out.write_text("\n".join(json.dumps(q) for q in queue) + "\n")120    origins = sorted({q["origin"] for q in queue})121    missing = sum(1 for q in queue if q["ctx_missing"])122    print(f"scanned={len(scanned)} queued={len(queue)} origins={origins} "123          f"ctx_missing={missing} -> {out}")124 125 126if __name__ == "__main__":127    main()128