spb/wpai0002_uqo

1from htr.data.splits import SplitConfig, assign_splits, check_no_group_leak, split_summary2from htr.dedup.minhash import Deduper345def test_splits_group_safe():6    docs = [{"doc_id": f"d{i}", "author_group": f"a{i % 7}", "venue": "v", "subject_area": "physics" if i % 10 == 0 else "bio", "pub_year": 1950 if i % 15 == 0 else 2010} for i in range(60)]7    clusters = {d["doc_id"]: d["doc_id"] for d in docs}8    clusters["d1"] = "d2"  # near-duplicate pair9    asg = assign_splits(docs, clusters, SplitConfig(heldout_disciplines=["physics"], heldout_before_year=1990))10    assert check_no_group_leak(asg) == []11    s = split_summary(asg)12    assert s.get("heldout_domain", 0) >= 1 and s.get("heldout_period", 0) >= 113    by = {a.doc_id: a.split for a in asg}14    assert by["d1"] == by["d2"] or {by["d1"], by["d2"]} & {"heldout_domain", "heldout_period"}15    # author groups never straddle splits16    for g in {d["author_group"] for d in docs}:17        splits = {by[d["doc_id"]] for d in docs if d["author_group"] == g} - {"heldout_domain", "heldout_period"}18        assert len(splits) <= 1192021def test_dedup():22    dd = Deduper(threshold=0.8)23    words = "alpha beta gamma delta epsilon zeta eta theta iota kappa lambda mu nu xi omicron pi rho sigma tau upsilon phi chi psi omega".split()24    base = " ".join(f"{words[i % 24]} {words[(i * 7) % 24]} {words[(i * 11) % 24]}" for i in range(80))25    assert dd.add("a", base) == "a"26    assert dd.add("b", base) == "a"  # exact27    assert dd.add("c", base.replace("alpha", "ALPHAX", 3)) == "a"  # near28    assert dd.add("d", "completely different content about economics and finance " * 20) == "d"29    assert "b" in dd.result.exact_dupes and "c" in dd.result.near_dupes30