"""CPU stage: country-agnostic pair features on rule-free text views.""" import math from multiprocessing import Pool import numpy as np import pandas as pd from rapidfuzz import fuzz, process from rapidfuzz.distance import JaroWinkler, Levenshtein import text_utils as tu from common import Progress, get_logger, n_workers, save_parquet, shard_path from prep import load_prep from retrieve_dense import block_rows, n_shards log = get_logger("feat") PCOLS = ["name_full", "name_core", "legal", "name_ns", "nonlatin", "addr_norm", "street", "street_ns", "loc", "house", "nums", "postal", "alnum", "addr_empty", "src"] PY_FEATS = ["n_soft_ab", "n_soft_ba", "n_jacc", "n_extra_a", "n_extra_b", "n_contain", "n_acro", "n_first_eq", "legal_state", "st_soft_ab", "st_soft_ba", "loc_cov_a", "loc_cov_b", "a_jacc", "house_state", "house_diff", "num_jacc", "num_state", "postal_state", "alnum_state", "n_tok_a", "n_tok_b", "a_tok_a", "a_tok_b"] FUZZ = [("n_ratio", "name_core", fuzz.ratio), ("n_jw", "name_core", JaroWinkler.normalized_similarity), ("n_lev", "name_core", Levenshtein.normalized_similarity), ("n_tsort", "name_core", fuzz.token_sort_ratio), ("n_tset", "name_core", fuzz.token_set_ratio), ("n_partial", "name_core", fuzz.partial_ratio), ("n_full_ratio", "name_full", fuzz.ratio), ("ns_ratio", "name_ns", fuzz.ratio), ("ns_partial", "name_ns", fuzz.partial_ratio), ("a_ratio", "addr_norm", fuzz.ratio), ("a_tset", "addr_norm", fuzz.token_set_ratio), ("a_partial", "addr_norm", fuzz.partial_ratio), ("st_ratio", "street", fuzz.ratio), ("st_ns_ratio", "street_ns", fuzz.ratio)] def _jacc(a, b): if not a and not b: return math.nan return len(a & b) / len(a | b) def _py_chunk(args): A, B = args # dicts of column -> list for the query side and candidate side out = np.empty((len(A["name_core"]), len(PY_FEATS)), dtype=np.float32) for i in range(len(out)): na, nb = A["name_core"][i].split(), B["name_core"][i].split() sa, sb = set(na), set(nb) fa, fb = A["name_full"][i].split(), B["name_full"][i].split() n_ab, n_ba = tu.soft_cov(na, nb), tu.soft_cov(nb, na) extra_a = sum(1 for x in na if x not in sb and not any(tu.tok_sim(x, y) > 0 for y in nb)) extra_b = sum(1 for x in nb if x not in sa and not any(tu.tok_sim(x, y) > 0 for y in na)) contain = 1.0 if tu.contains_seq(na, fb) else (2.0 if tu.contains_seq(nb, fa) else 0.0) acro = 1.0 if (tu.acronym_match(na, nb) or tu.acronym_match(nb, na)) else 0.0 first = 1.0 if na and nb and na[0] == nb[0] else 0.0 leg = tu.legal_state(A["legal"][i], B["legal"][i]) sta, stb = A["street"][i].split(), B["street"][i].split() st_ab, st_ba = tu.soft_cov(sta, stb), tu.soft_cov(stb, sta) ta, tb = A["addr_norm"][i].split(), B["addr_norm"][i].split() la = [t for seg in A["loc"][i].split("|") for t in seg.split() if not t.isdigit()] lb = [t for seg in B["loc"][i].split("|") for t in seg.split() if not t.isdigit()] loc_a = tu.soft_cov(la, tb) # does the query's city/region appear anywhere in the candidate loc_b = tu.soft_cov(lb, ta) hs, hd = tu.house_compare(A["house"][i], B["house"][i]) ma, mb = set(A["nums"][i].split()), set(B["nums"][i].split()) pa, pb = set(A["postal"][i].split()), set(B["postal"][i].split()) xa, xb = set(A["alnum"][i].split()), set(B["alnum"][i].split()) out[i] = (n_ab, n_ba, _jacc(sa, sb), extra_a, extra_b, contain, acro, first, leg, st_ab, st_ba, loc_a, loc_b, _jacc(set(ta), set(tb)), hs, hd, _jacc(ma, mb), tu.set_state(ma, mb), tu.set_state(pa, pb), tu.set_state(xa, xb), len(na), len(nb), len(ta), len(tb)) return out def _gather(tab, idx, cols): return {c: tab[c][idx].tolist() for c in cols} def pair_features(pool, s1tab, pooltab, df): q, c = df.q.values, df.c.values feats = {} for name, col, scorer in FUZZ: a, b = s1tab[col][q].tolist(), pooltab[col][c].tolist() # lists: ~10x faster than object arrays heavy = scorer in (fuzz.partial_ratio, fuzz.token_set_ratio, fuzz.token_sort_ratio) v = process.cpdist(a, b, scorer=scorer, workers=16 if heavy else 1).astype(np.float32) empty = (pd.Series(a).str.len().values == 0) | (pd.Series(b).str.len().values == 0) v[empty] = np.nan feats[name] = v str_cols = ["name_core", "name_full", "legal", "street", "addr_norm", "loc", "house", "nums", "postal", "alnum"] step = 20_000 args = [(_gather(s1tab, q[i:i + step], str_cols), _gather(pooltab, c[i:i + step], str_cols)) for i in range(0, len(df), step)] py = np.concatenate(list(pool.imap(_py_chunk, args))) if args else np.empty((0, len(PY_FEATS)), np.float32) for j, name in enumerate(PY_FEATS): feats[name] = py[:, j] feats["nonlatin_a"] = s1tab["nonlatin"][q].astype(np.float32) feats["nonlatin_b"] = pooltab["nonlatin"][c].astype(np.float32) feats["addr_empty_a"] = s1tab["addr_empty"][q].astype(np.float32) feats["addr_empty_b"] = pooltab["addr_empty"][c].astype(np.float32) feats["cand_src"] = pooltab["src"][c].astype(np.float32) return pd.DataFrame(feats) def query_context(df): g = df.groupby("q") out = pd.DataFrame(index=df.index) out["q_ncand"] = g.c.transform("size").astype(np.float32) for m in ("e5", "bge"): best = g[f"{m}_cos"].transform("max") out[f"{m}_cos_gap"] = (best - df[f"{m}_cos"]).astype(np.float32) out[f"q_best_{m}"] = best.astype(np.float32) out[f"{m}_cos_rank"] = g[f"{m}_cos"].rank(ascending=False, method="first").astype(np.float32) out["q_best_n_tset"] = g.n_tset.transform("max").astype(np.float32) return out def run_features(split): ns = n_shards(split) todo = [i for i in range(ns) if not shard_path("feat", split, i).exists()] if not todo: return with Pool(n_workers()) as pool: # fork before big tables are loaded s1 = load_prep(split, "s1", [c for c in PCOLS if c != "src"]) pl = load_prep(split, "pool", PCOLS) s1tab = {c: s1[c].values for c in s1.columns} pooltab = {c: pl[c].values for c in pl.columns} del s1, pl prog = Progress(log, len(todo), f"features {split} shards") for i in todo: df = block_rows(split, i) for m in ("e5", "bge"): z = np.load(shard_path("cos", split, i, "npz", m)) for k in z.files: df[k] = z[k] f = pair_features(pool, s1tab, pooltab, df) df = pd.concat([df, f], axis=1) df = pd.concat([df, query_context(df)], axis=1) save_parquet(df, shard_path("feat", split, i)) prog.step() def feature_columns(df): return [c for c in df.columns if c not in ("q", "c", "y")]