| """Loading, S1 augmentation with ground truth, sanity checks and splits. |
| |
| Unchanged from model-2 except: load_data() takes an explicit paths dict so the |
| same code serves train and test, and test loading needs no ground truth. |
| """ |
| import os |
| import numpy as np |
| import pandas as pd |
|
|
| S1_COLS = ["entity_id", "business_name", "business_address", "country", "country_norm", |
| "name_norm", "name_core", "legal_suffix", "address_abbrev_norm", "address_numbers"] |
| S23_COLS = S1_COLS |
|
|
|
|
| def read_any(path, columns=None): |
| if path.endswith(".parquet"): |
| df = pd.read_parquet(path, columns=columns) |
| else: |
| df = pd.read_csv(path, sep="\t", dtype=str, keep_default_na=False, usecols=columns) |
| for c in df.columns: |
| df[c] = df[c].fillna("").astype(str) |
| return df |
|
|
|
|
| def split_ids(s): |
| return [x for x in str(s).split(",") if x] |
|
|
|
|
| def load_data(paths, with_gt=True): |
| s1 = read_any(paths["source1"], S1_COLS) |
| s2 = read_any(paths["source2"], S23_COLS) |
| s3 = read_any(paths["source3"], S23_COLS) |
| s23 = pd.concat([s2, s3], ignore_index=True) |
| if not with_gt: |
| return s1, s23, None |
| gt = read_any(paths["ground_truth"]) |
| return s1, s23, gt |
|
|
|
|
| def augment_s1(s1, gt, out_path=None): |
| """Attach ground-truth labels to Source 1 as column `matched_entity_ids`.""" |
| assert s1["entity_id"].is_unique, "duplicate entity_id in source1" |
| assert gt["source1_entity_id"].is_unique, "duplicate source1_entity_id in ground truth" |
| s1 = s1.merge(gt.rename(columns={"source1_entity_id": "entity_id"}), on="entity_id", |
| how="left", validate="one_to_one") |
| missing = (s1["matched_entity_ids"].isna()).sum() |
| assert missing == 0, f"{missing} source1 entities have no ground-truth row" |
| s1["matched_entity_ids"] = s1["matched_entity_ids"].fillna("").astype(str) |
| if out_path: |
| s1.to_parquet(out_path, index=False) |
| return s1 |
|
|
|
|
| def sanity_checks(s1, s23, rng, sample=200_000): |
| rep = {} |
| rep["n_s1"] = len(s1) |
| rep["n_s23"] = len(s23) |
| rep["s23_ids_unique"] = bool(s23["entity_id"].is_unique) |
| lists = s1["matched_entity_ids"].map(split_ids) |
| n = lists.map(len) |
| rep["singleton_rate"] = float((n == 0).mean()) |
| rep["mean_matches"] = float(n.mean()) |
| rep["max_matches"] = int(n.max()) |
| flat = pd.Series([x for l in lists for x in l]) |
| rep["n_matched_ids"] = int(len(flat)) |
| rep["multi_owner_ids"] = int(flat.duplicated().sum()) |
| idx = pd.Index(s23["entity_id"]) |
| pos = idx.get_indexer(flat) |
| rep["gt_ids_missing_in_s23"] = int((pos < 0).sum()) |
| |
| owner = np.repeat(np.arange(len(s1)), n.values) |
| take = rng.choice(len(flat), size=min(sample, len(flat)), replace=False) |
| ok = pos[take] >= 0 |
| q_c = s1["country_norm"].values[owner[take][ok]] |
| c_c = s23["country_norm"].values[pos[take][ok]] |
| rep["true_pairs_same_country_rate"] = float((q_c == c_c).mean()) if ok.any() else None |
| rep["s1_country_counts"] = s1["country_norm"].value_counts().to_dict() |
| return rep |
|
|
|
|
| def make_split(s1, cfg, rng): |
| """Return three disjoint arrays of S1 row positions: ce, stack, val.""" |
| n_ce, n_st, n_va = cfg["n_s1_ce"], cfg["n_s1_stack"], cfg["n_s1_val"] |
| if cfg["split"] == "A": |
| perm = rng.permutation(len(s1)) |
| need = n_ce + n_st + n_va |
| if need > len(perm): |
| raise ValueError(f"requested {need} S1 entities but only {len(perm)} exist") |
| return perm[:n_ce], perm[n_ce:n_ce + n_st], perm[n_ce + n_st:need] |
| if cfg["split"] == "B": |
| c = s1["country_norm"].values |
| tr = rng.permutation(np.where(c == cfg["split_b_train_country"])[0]) |
| va = rng.permutation(np.where(c == cfg["split_b_val_country"])[0]) |
| if n_ce + n_st > len(tr) or n_va > len(va): |
| raise ValueError("not enough entities for split B with these sizes") |
| return tr[:n_ce], tr[n_ce:n_ce + n_st], va[:n_va] |
| raise ValueError(cfg["split"]) |
|
|
|
|
| def ensure_dir(p): |
| os.makedirs(p, exist_ok=True) |
| return p |
|
|
