File size: 7,008 Bytes
c43c544 | 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 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """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")]
|