| """Pair features for the LightGBM scorer. All features are country-agnostic.""" |
| import multiprocessing as mp |
| import numpy as np |
| import pandas as pd |
| from rapidfuzz import process, fuzz |
| from rapidfuzz.distance import JaroWinkler, Levenshtein |
|
|
| from text_utils import abbrev_aware_match, token_jaccard, numbers |
|
|
|
|
| def _cp(a, b, scorer, workers): |
| return process.cpdist(a, b, scorer=scorer, workers=workers).astype(np.float32) |
|
|
|
|
| _A = _B = None |
|
|
|
|
| def _py_feats(rng): |
| lo, hi = rng |
| out = np.empty((hi - lo, 6), np.float32) |
| for j, i in enumerate(range(lo, hi)): |
| a, b = _A[i], _B[i] |
| out[j, 0] = abbrev_aware_match(a["name"], b["name"]) |
| out[j, 1] = abbrev_aware_match(a["addr"], b["addr"]) |
| out[j, 2] = token_jaccard(a["name"], b["name"]) |
| out[j, 3] = token_jaccard(a["addr"], b["addr"]) |
| na, nb = numbers(a["nums"]), numbers(b["nums"]) |
| out[j, 4] = len(na & nb) / len(na | nb) if (na or nb) else np.nan |
| la = {x for x in na if len(x) >= 4} |
| lb = {x for x in nb if len(x) >= 4} |
| out[j, 5] = (1.0 if la & lb else 0.0) if (la and lb) else np.nan |
| return out |
|
|
|
|
| def compute_features(cand, s1v, s23v, n_jobs=8): |
| """cand: DataFrame with q (S1 row), c (S23 row) and blocking columns. |
| s1v / s23v: dict of numpy arrays per column (name_core_abbr, address_abbrev_norm, ...).""" |
| q, c = cand["q"].values, cand["c"].values |
| f = pd.DataFrame(index=cand.index) |
| na, nb = s1v["name_core_abbr"][q], s23v["name_core_abbr"][c] |
| fa, fb = s1v["name_norm"][q], s23v["name_norm"][c] |
| aa, ab = s1v["address_abbrev_norm"][q], s23v["address_abbrev_norm"][c] |
| w = n_jobs |
|
|
| f["name_jw"] = _cp(na, nb, JaroWinkler.normalized_similarity, w) |
| f["name_lev"] = _cp(na, nb, Levenshtein.normalized_similarity, w) |
| f["name_tsort"] = _cp(na, nb, fuzz.token_sort_ratio, w) / 100 |
| f["name_tset"] = _cp(na, nb, fuzz.token_set_ratio, w) / 100 |
| f["name_partial"] = _cp(na, nb, fuzz.partial_ratio, w) / 100 |
| f["fullname_ratio"] = _cp(fa, fb, fuzz.ratio, w) / 100 |
| f["addr_ratio"] = _cp(aa, ab, fuzz.ratio, w) / 100 |
| f["addr_tset"] = _cp(aa, ab, fuzz.token_set_ratio, w) / 100 |
| f["addr_partial"] = _cp(aa, ab, fuzz.partial_ratio, w) / 100 |
|
|
| global _A, _B |
| _A = [{"name": x, "addr": y, "nums": z} for x, y, z in |
| zip(na, aa, s1v["nums"][q])] |
| _B = [{"name": x, "addr": y, "nums": z} for x, y, z in |
| zip(nb, ab, s23v["nums"][c])] |
| n = len(cand) |
| step = max(1, n // (n_jobs * 8) + 1) |
| ranges = [(i, min(i + step, n)) for i in range(0, n, step)] |
| if n_jobs > 1 and n > 50_000 and "fork" in mp.get_all_start_methods(): |
| with mp.get_context("fork").Pool(n_jobs) as pool: |
| parts = pool.map(_py_feats, ranges) |
| else: |
| parts = [_py_feats(r) for r in ranges] |
| pf = np.concatenate(parts) if parts else np.empty((0, 6), np.float32) |
| for j, name in enumerate(["name_abbrev", "addr_abbrev", "name_jacc", "addr_jacc", |
| "num_jacc", "longnum_match"]): |
| f[name] = pf[:, j] |
| _A = _B = None |
|
|
| la, lb = s1v["legal_suffix"][q], s23v["legal_suffix"][c] |
| f["suffix_both"] = ((la != "") & (lb != "")).astype(np.int8) |
| f["suffix_eq"] = np.where(f["suffix_both"] == 1, (la == lb).astype(np.float32), np.nan) |
| f["same_country"] = (s1v["country_norm"][q] == s23v["country_norm"][c]).astype(np.int8) |
| f["name_len_a"] = np.fromiter((len(x) for x in na), np.int32, n) |
| f["name_len_b"] = np.fromiter((len(x) for x in nb), np.int32, n) |
| f["addr_empty_any"] = ((aa == "") | (ab == "")).astype(np.int8) |
| f["is_s3"] = np.char.startswith(s23v["entity_id"][c].astype(str), "S3-").astype(np.int8) |
|
|
| |
| |
| g = cand.groupby("q") |
| f["emb_gap"] = cand["emb_score"].values - g["emb_score"].transform("max").values |
| f["lex_gap"] = cand["lex_score"].values - g["lex_score"].transform("max").values |
| f["n_cands"] = g["c"].transform("size").values |
| f["name_jw_gap"] = f["name_jw"].values - f.groupby(cand["q"].values)["name_jw"].transform("max").values |
| return f |
|
|
