File size: 4,183 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 | """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)
# blocking meta + within-query context
# (raw blocking columns emb_score/emb_rank/lex_score/lex_rank/n_keys stay in `cand`)
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
|