Repository Public Research preview
1"""Leakage-safe splits. Unit of assignment = near-duplicate cluster of documents; groups = author group and2venue so that no author/venue group straddles train and test when ``author_safe``/``venue_safe`` are set.3Held-out discipline(s) and the earliest period can be carved out entirely. Deterministic given seed."""4from __future__ import annotations56import hashlib7import random8from collections import defaultdict9from dataclasses import dataclass, field1011from htr.data.schema import SplitAssignment121314@dataclass15class SplitConfig:16 seed: int = 2026110917 train: float = 0.818 val: float = 0.119 test: float = 0.120 heldout_disciplines: list[str] = field(default_factory=list)21 heldout_before_year: int | None = None # documents with pub_year < this go to heldout_period22 author_safe: bool = True23 venue_safe: bool = False242526def _stable_rand(key: str, seed: int) -> float:27 h = hashlib.sha256(f"{seed}:{key}".encode()).hexdigest()28 return int(h[:12], 16) / 16**12293031def assign_splits(docs: list[dict], clusters: dict[str, str], cfg: SplitConfig, fixed: dict[str, str] | None = None) -> list[SplitAssignment]:32 """``docs``: dicts with doc_id, author_group, venue, subject_area, pub_year. ``clusters``: doc_id → cluster id.33 Union-find over cluster + author group (+ venue) builds leakage groups; groups are then hashed into splits.34 ``fixed`` = previously published assignments (doc_id → split) that must not change (passages already rewritten);35 any group containing a fixed document inherits that split (conflicts → 'quarantine')."""36 parent: dict[str, str] = {}3738 def find(x: str) -> str:39 while parent.setdefault(x, x) != x:40 parent[x] = parent[parent[x]]41 x = parent[x]42 return x4344 def union(a: str, b: str) -> None:45 ra, rb = find(a), find(b)46 if ra != rb:47 parent[rb] = ra4849 for d in docs:50 did = d["doc_id"]51 union(f"doc:{did}", f"clu:{clusters.get(did, did)}")52 if cfg.author_safe and d.get("author_group"):53 union(f"doc:{did}", f"auth:{d['author_group']}")54 if cfg.venue_safe and d.get("venue"):55 union(f"doc:{did}", f"venue:{d['venue']}")56 fixed = fixed or {}57 group_fixed: dict[str, set[str]] = defaultdict(set)58 for did, sp in fixed.items():59 group_fixed[find(f"doc:{did}")].add(sp)60 out: list[SplitAssignment] = []61 for d in docs:62 did = d["doc_id"]63 g = find(f"doc:{did}")64 if g in group_fixed:65 fs = group_fixed[g]66 sp = next(iter(fs)) if len(fs) == 1 else "quarantine"67 out.append(SplitAssignment(doc_id=did, split=sp, group_keys={"group": g, "cluster": clusters.get(did, did), "author_group": d.get("author_group") or "", "venue": d.get("venue") or "", "fixed": "1"}))68 continue69 keys = {"group": g, "cluster": clusters.get(did, did), "author_group": d.get("author_group") or "", "venue": d.get("venue") or ""}70 if d.get("subject_area") in cfg.heldout_disciplines:71 out.append(SplitAssignment(doc_id=did, split="heldout_domain", group_keys=keys))72 continue73 if cfg.heldout_before_year and d.get("pub_year") and d["pub_year"] < cfg.heldout_before_year:74 out.append(SplitAssignment(doc_id=did, split="heldout_period", group_keys=keys))75 continue76 r = _stable_rand(g, cfg.seed)77 split = "train" if r < cfg.train else ("val" if r < cfg.train + cfg.val else "test")78 out.append(SplitAssignment(doc_id=did, split=split, group_keys=keys))79 return out808182def split_summary(assignments: list[SplitAssignment]) -> dict[str, int]:83 c: dict[str, int] = defaultdict(int)84 for a in assignments:85 c[a.split] += 186 return dict(sorted(c.items()))878889def check_no_group_leak(assignments: list[SplitAssignment]) -> list[str]:90 seen: dict[str, set[str]] = defaultdict(set)91 for a in assignments:92 seen[a.group_keys["group"]].add(a.split)93 return [g for g, s in seen.items() if len(s - {"heldout_domain", "heldout_period"}) > 1]949596def shuffled(items: list, seed: int) -> list:97 r = random.Random(seed)98 items = list(items)99 r.shuffle(items)100 return items101