File size: 4,052 Bytes
e79d917 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """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())
# same-country rate over a sample of true pairs
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
|