| """Stages: fetch (Hugging Face inputs), manifest (integrity), prep (normalised views + legal forms), test_tsv, |
| splits.""" |
| import os |
| import random |
| from collections import Counter |
| from multiprocessing import Pool |
| from pathlib import Path |
|
|
| import numpy as np |
| import pandas as pd |
| import pyarrow.parquet as pq |
|
|
| import text_utils as tu |
| from common import Progress, get_logger, load_json, n_workers, save_json, save_npz, save_parquet |
| from config import CFG, EMB, EMB_MODELS, HF, PATHS, PREP, WORK |
| from embeddings import emb_ids, load_emb |
|
|
| log = get_logger("prep") |
| RAW_COLS = ["entity_id", "business_name", "business_address", "country"] |
| SIDES = [("train", "s1"), ("train", "pool"), ("test", "s1"), ("test", "pool")] |
|
|
|
|
| |
|
|
| def _needed_files(): |
| need = {"data": [p for sp in PATHS["data"].values() for p in sp.values()] + [PATHS["ground_truth"]]} |
| for m in EMB_MODELS: |
| spec = EMB[m] |
| fs = [] |
| for split in ("train", "test"): |
| d = Path(spec[split]) |
| if spec["kind"] == "shards": |
| fs += [d / f"source{s}" / "manifest.json" for s in (1, 2, 3)] |
| elif spec["kind"] == "s1pool": |
| fs += [d / f"{s}{x}" for s in ("s1", "pool") for x in (".npy", ".ids.txt")] |
| else: |
| fs += [d / f"{split}_{s}" / x for s in ("s1", "pool") for x in ("matrix.npy", "ids.txt")] |
| need[m] = fs |
| return need |
|
|
|
|
| def run_fetch(): |
| """Download any missing input from Hugging Face (snapshot_download resumes partial files). Nothing is |
| downloaded when every file is already present.""" |
| missing = {k: [str(f) for f in fs if not Path(f).exists()] for k, fs in _needed_files().items()} |
| missing = {k: v for k, v in missing.items() if v} |
| if not missing: |
| log.info("all inputs present locally; nothing to download") |
| return |
| from huggingface_hub import snapshot_download |
| token = os.environ.get("HF_TOKEN") |
| for key, files in missing.items(): |
| repo, local, allow = HF[key] |
| log.info(f"{key}: {len(files)} missing file(s), e.g. {files[0]} -> downloading {repo} into {local}") |
| for attempt in range(5): |
| try: |
| snapshot_download(repo_id=repo, repo_type="dataset", local_dir=str(local), token=token, |
| allow_patterns=allow, max_workers=8) |
| break |
| except Exception as e: |
| log.warning(f"{repo}: download attempt {attempt + 1} failed: {e}") |
| if attempt == 4: |
| raise |
| import time |
| time.sleep(60 * (attempt + 1)) |
| still = [str(f) for fs in _needed_files().values() for f in fs if not Path(f).exists()] |
| if still: |
| raise RuntimeError(f"inputs still missing after download: {still[:5]}") |
|
|
|
|
| |
|
|
| def read_raw(split, side): |
| files = [PATHS["data"][split][1]] if side == "s1" else [PATHS["data"][split][2], PATHS["data"][split][3]] |
| df = pd.concat([pq.read_table(f, columns=RAW_COLS).to_pandas() for f in files], ignore_index=True) |
| for c in RAW_COLS: |
| df[c] = df[c].fillna("").astype(str) |
| return df |
|
|
|
|
| def prep_path(split, side): |
| return PREP / f"{split}_{side}.parquet" |
|
|
|
|
| def load_prep(split, side, columns=None): |
| return pd.read_parquet(prep_path(split, side), columns=columns) |
|
|
|
|
| def _norm_scan(e, rng, n=50_000): |
| rows = np.unique(np.concatenate([rng.choice(len(e), min(n, len(e)), replace=False), |
| np.arange(min(1000, len(e))), np.arange(max(0, len(e) - 1000), len(e))])) |
| x = e.take(rows).astype(np.float32) |
| nr = np.linalg.norm(x, axis=1) |
| return {"sampled": int(len(rows)), "norm_min": float(np.nanmin(nr)), "norm_max": float(np.nanmax(nr)), |
| "norm_bad": int((~np.isfinite(nr) | (np.abs(nr - 1) > 0.01)).sum())} |
|
|
|
|
| def run_manifest(): |
| """Fatal integrity checks: row counts, id order (id_mismatch must be 0), unit norm, and a semantic row |
| alignment test on ground-truth pairs (true pair vs the pair shifted by one row).""" |
| rep, problems = {}, [] |
| ids = {} |
| rng = np.random.default_rng(0) |
| for split, side in SIDES: |
| df = read_raw(split, side) |
| n = len(df) |
| dup = int(df.entity_id.duplicated().sum()) |
| ok_pref = df.entity_id.str[:3].isin(["S1-"] if side == "s1" else ["S2-", "S3-"]).all() |
| r = rep[f"{split}_{side}"] = {"rows": n, "dup_ids": dup, "prefix_ok": bool(ok_pref), |
| "countries": df.country.value_counts().to_dict(), |
| "empty_name": int((df.business_name == "").sum()), |
| "empty_addr": int((df.business_address == "").sum()), "emb": {}} |
| if dup or not ok_pref: |
| problems.append(f"{split}_{side}: dup={dup} prefix_ok={ok_pref}") |
| ids[(split, side)] = df.entity_id.values |
| for m in EMB_MODELS: |
| e = load_emb(m, split, side) |
| er = r["emb"][m] = {"rows": len(e), "dim": e.dim, **_norm_scan(e, rng)} |
| if len(e) != n: |
| problems.append(f"{m} {split}_{side}: {len(e)} emb rows vs {n} data rows") |
| eids = emb_ids(m, split, side) |
| if eids is not None: |
| mm = int((np.asarray(eids, dtype=object) != df.entity_id.values).sum()) if len(eids) == n else None |
| er["id_mismatch"] = mm |
| if mm != 0: |
| problems.append(f"{m} {split}_{side}: id_mismatch={mm}") |
| if er["norm_bad"]: |
| problems.append(f"{m} {split}_{side}: {er['norm_bad']} sampled rows not unit-norm") |
| log.info(f"checked {split}_{side}: {r}") |
| gt = pd.read_csv(PATHS["ground_truth"], sep="\t", dtype=str, keep_default_na=False) |
| s1ids, pool_ids = ids[("train", "s1")], ids[("train", "pool")] |
| matched = [x for m in gt.matched_entity_ids if m for x in m.split(",")] |
| cnt = Counter(matched) |
| pool_set = set(pool_ids) |
| rep["ground_truth"] = {"rows": len(gt), "s1_covered": int(gt.source1_entity_id.isin(set(s1ids)).sum()), |
| "n_matched": len(matched), "unknown_ids": int(sum(1 for x in cnt if x not in pool_set)), |
| "multi_owner": int(sum(1 for v in cnt.values() if v > 1))} |
| if rep["ground_truth"]["unknown_ids"] or rep["ground_truth"]["s1_covered"] != len(s1ids): |
| problems.append(f"ground truth coverage: {rep['ground_truth']}") |
| if rep["ground_truth"]["multi_owner"]: |
| problems.append("S2/S3 ids owned by >1 S1: one-owner arbitration invalid") |
| s1_index = pd.Index(s1ids) |
| pool_index = pd.Index(pool_ids) |
| sub = gt[gt.matched_entity_ids != ""].sample(20_000, random_state=0) |
| a = s1_index.get_indexer(sub.source1_entity_id) |
| b = pool_index.get_indexer(sub.matched_entity_ids.str.split(",").str[0]) |
| ok = (a + 1 < len(s1ids)) & (b + 1 < len(pool_ids)) |
| a, b = a[ok], b[ok] |
| rand = rng.choice(len(pool_ids), len(b)) |
| for m in EMB_MODELS: |
| es, ep = load_emb(m, "train", "s1"), load_emb(m, "train", "pool") |
| qa, pb = es.take(a).astype(np.float32), ep.take(b).astype(np.float32) |
| t = (qa * pb).sum(1) |
| sh_q = (es.take(a + 1).astype(np.float32) * pb).sum(1) |
| sh_p = (qa * ep.take(b + 1).astype(np.float32)).sum(1) |
| rnd = (qa * ep.take(rand).astype(np.float32)).sum(1) |
| al = {"cos_true": float(t.mean()), "cos_random": float(rnd.mean()), |
| "frac_true_beats_s1_shift": float((t > sh_q).mean()), "frac_true_beats_pool_shift": float((t > sh_p).mean())} |
| rep["ground_truth"][f"{m}_alignment"] = al |
| if al["frac_true_beats_s1_shift"] < 0.95 or al["frac_true_beats_pool_shift"] < 0.95 or \ |
| al["cos_true"] - al["cos_random"] < 0.1: |
| problems.append(f"{m}: embedding rows not aligned with parquet rows: {al}") |
| rep["problems"] = problems |
| save_json(rep, WORK / "reports" / "manifest.json") |
| if problems: |
| raise RuntimeError("integrity manifest failed: " + "; ".join(problems)) |
| log.info("manifest OK") |
|
|
|
|
| |
|
|
| _FORMS = None |
|
|
|
|
| def _init(forms): |
| global _FORMS |
| _FORMS = forms |
|
|
|
|
| def _tok_chunk(names): |
| return [tu.name_tokens(x) for x in names] |
|
|
|
|
| def _proc_chunk(args): |
| names, addrs, countries = args |
| rows = [] |
| for n, a, c in zip(names, addrs, countries): |
| forms, long_single = _FORMS.get(c, ([], set())) |
| r = tu.process_name(n, forms, long_single) |
| r.update(tu.parse_address(a)) |
| rows.append(r) |
| return pd.DataFrame(rows) |
|
|
|
|
| def norm_country(s): |
| return s.strip().lower() or "unknown" |
|
|
|
|
| def run_prep(): |
| PREP.mkdir(parents=True, exist_ok=True) |
| forms_path = PREP / "legal_forms.json" |
| raws = {} |
| with Pool(n_workers()) as pool: |
| if not forms_path.exists(): |
| by_c = {} |
| for split, side in SIDES: |
| df = read_raw(split, side) |
| df["cn"] = df.country.map(norm_country) |
| raws[(split, side)] = df |
| for c, g in df.groupby("cn"): |
| by_c.setdefault(c, []).extend(g.business_name.tolist()) |
| rng = random.Random(CFG["seed"]) |
| forms = {} |
| for c, names in by_c.items(): |
| if len(names) > CFG["legal_sample"]: |
| names = rng.sample(names, CFG["legal_sample"]) |
| chunks = [names[i:i + 50_000] for i in range(0, len(names), 50_000)] |
| toks = [t for part in pool.map(_tok_chunk, chunks) for t in part] |
| f = tu.discover_legal_forms(toks, CFG["legal_min_frac"], CFG["legal_boundary_ratio"], |
| CFG["legal_max_cooc"], CFG["legal_bigram_tight"]) |
| forms[c] = [" ".join(x) for x in f] |
| log.info(f"legal forms [{c}] from {len(names)} names: {forms[c]}") |
| save_json(forms, forms_path) |
| forms = load_json(forms_path) |
| forms_t = {c: ([tuple(x.split()) for x in v], {x for x in v if " " not in x and len(x) >= 3}) |
| for c, v in forms.items()} |
| with Pool(n_workers(), initializer=_init, initargs=(forms_t,)) as pool: |
| for split, side in SIDES: |
| out = prep_path(split, side) |
| if out.exists(): |
| continue |
| df = raws.pop((split, side), None) |
| if df is None: |
| df = read_raw(split, side) |
| df["cn"] = df.country.map(norm_country) |
| step = 100_000 |
| args = [(df.business_name.values[i:i + step], df.business_address.values[i:i + step], |
| df.cn.values[i:i + step]) for i in range(0, len(df), step)] |
| prog = Progress(log, len(args), f"prep {split}_{side}") |
| parts = [] |
| for part in pool.imap(_proc_chunk, args): |
| parts.append(part) |
| prog.step() |
| feat = pd.concat(parts, ignore_index=True) |
| res = pd.DataFrame({"id": df.entity_id.values, "country": df.cn.values, |
| "src": df.entity_id.str[1].astype(np.int8).values, |
| "raw_name": df.business_name.values, "raw_addr": df.business_address.values, |
| "raw_country": df.country.values}) |
| res = pd.concat([res, feat], axis=1) |
| save_parquet(res, out) |
| log.info(f"wrote {out} rows={len(res)}") |
|
|
|
|
| def run_test_tsv(): |
| """The organisers' validator reads test_source{1,2,3}.tsv; write them from the parquet (same ids/rows).""" |
| d = PATHS["test_tsv"] |
| for s in (1, 2, 3): |
| out = d / f"test_source{s}.tsv" |
| if out.exists(): |
| continue |
| df = pq.read_table(PATHS["data"]["test"][s], columns=RAW_COLS).to_pandas() |
| tmp = out.with_name(f".{out.name}.tmp") |
| tmp.parent.mkdir(parents=True, exist_ok=True) |
| df.to_csv(tmp, sep="\t", index=False) |
| os.replace(tmp, out) |
| log.info(f"wrote {out} rows={len(df)}") |
|
|
|
|
| |
|
|
| GROUPS = ("ce", "lce", "stack", "val") |
|
|
|
|
| def run_splits(): |
| s1 = load_prep("train", "s1", ["id", "country"]) |
| pool_ids = load_prep("train", "pool", ["id"]).id.values |
| rng = np.random.default_rng(CFG["seed"]) |
| perm = rng.permutation(len(s1)) |
| |
| |
| sizes = [("ce", CFG["n_ce"]), ("stack", CFG["n_stack"]), ("val", CFG["n_val"]), ("lce", CFG["n_lce"])] |
| groups, o = {}, 0 |
| for g, n in sizes: |
| groups[g] = np.sort(perm[o:o + n]) |
| o += n |
| for a in GROUPS: |
| for b in GROUPS: |
| if a < b: |
| assert len(np.intersect1d(groups[a], groups[b])) == 0, f"split groups {a} and {b} overlap" |
| q_train = np.sort(np.concatenate(list(groups.values()))) |
| n_test = len(load_prep("test", "s1", ["id"])) |
| q_test = np.arange(n_test) |
| if CFG["test_limit"]: |
| q_test = np.sort(rng.choice(n_test, CFG["test_limit"], replace=False)) |
| gt = pd.read_csv(PATHS["ground_truth"], sep="\t", dtype=str, keep_default_na=False) |
| s1_index = pd.Series(np.arange(len(s1)), index=s1.id.values) |
| pool_index = pd.Series(np.arange(len(pool_ids)), index=pool_ids) |
| gt = gt[gt.matched_entity_ids != ""] |
| ex = gt.assign(m=gt.matched_entity_ids.str.split(",")).explode("m") |
| gq = s1_index.loc[ex.source1_entity_id.values].values.astype(np.int64) |
| gc = pool_index.loc[ex.m.values].values.astype(np.int64) |
| n_true = np.bincount(gq, minlength=len(s1)).astype(np.int16) |
| save_npz(WORK / "splits.npz", q_train=q_train, q_test=q_test, n_true=n_true, gt_q=gq, gt_c=gc, |
| **{f"g_{k}": v for k, v in groups.items()}) |
| rep = {k: {"n": int(len(v)), "countries": pd.Series(s1.country.values[v]).value_counts().to_dict(), |
| "singleton_rate": float((n_true[v] == 0).mean())} for k, v in groups.items()} |
| rep["q_test"] = int(len(q_test)) |
| save_json(rep, WORK / "reports" / "splits.json") |
| log.info(f"splits: {rep}") |
|
|
|
|
| def load_splits(): |
| z = np.load(WORK / "splits.npz") |
| return {k: z[k] for k in z.files} |
|
|
