"""CPU stage: country-agnostic pair features on rule-free text views, plus country-invariant versions of the scale-dependent ones (README F1).""" import math import sys import time 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, load_json, n_workers, save_parquet, shard_path, claim, release from config import EMB_MODELS, WORK 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", "country"] 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)] # features whose scale depends on a country's vocabulary / pool / language (README F1). The F1 selection stage # decides under leave-one-country-out whether to drop them; their invariant replacements are always present. def raw_scale_features(cols): raw = {f"{r}_{x}" for r in ("e5", "bge", "e5s", "lex", "hop") for x in ("s", "gap")} raw |= {"e5_cos", "bge_cos", "e5s_cos", "q_best_e5", "q_best_bge", "q_best_e5s", "tri_max_cos", "n_tok_a", "n_tok_b", "a_tok_a", "a_tok_b"} return sorted(c for c in cols if c in raw) 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, nthreads): 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 v = process.cpdist(a, b, scorer=scorer, workers=nthreads).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 = 10_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, retr, stats, s1_country): """Within-query ranks / gaps / z-scores, and per-country standardised cosines (label-free statistics).""" g = df.groupby("q") out = pd.DataFrame(index=df.index) out["q_ncand"] = g.c.transform("size").astype(np.float32) ctry = s1_country[df.q.values] for m in EMB_MODELS: col = f"{m}_cos" best = g[col].transform("max") mu_q, sd_q = g[col].transform("mean"), g[col].transform("std").fillna(0) out[f"{m}_cos_gap"] = (best - df[col]).astype(np.float32) out[f"q_best_{m}"] = best.astype(np.float32) out[f"{m}_cos_rank"] = g[col].rank(ascending=False, method="first").astype(np.float32) out[f"{m}_cos_qz"] = ((df[col] - mu_q) / (sd_q + 1e-3)).astype(np.float32) st = stats[m] mu = np.array([st.get(c, {"mu": np.nan})["mu"] for c in ctry], dtype=np.float32) sd = np.array([st.get(c, {"sd": np.nan})["sd"] for c in ctry], dtype=np.float32) out[f"{m}_cos_cz"] = ((df[col].values - mu) / sd).astype(np.float32) out[f"q_best_{m}_cz"] = ((best.values - mu) / sd).astype(np.float32) if "lex_s" in df: mx = g.lex_s.transform("max") out["lex_rel"] = (df.lex_s / mx.where(mx > 0)).astype(np.float32) out["q_best_n_tset"] = g.n_tset.transform("max").astype(np.float32) out["n_tset_gap"] = (out.q_best_n_tset - df.n_tset).astype(np.float32) out["tok_ratio_n"] = (df.n_tok_a / df.n_tok_b.clip(lower=1)).astype(np.float32) out["tok_ratio_a"] = (df.a_tok_a / df.a_tok_b.clip(lower=1)).astype(np.float32) return out def run_features(split): ns = n_shards(split) paths = [shard_path("feat", split, i) for i in range(ns)] if all(p.exists() for p in paths): return from config import RETR # noqa: F401 (the kept retrievers are already in the block shards) stats = {m: load_json(WORK / "reports" / f"cos_stats_{m}_{split}.json") for m in EMB_MODELS} nw = n_workers() with Pool(nw) 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} s1_country = s1.country.values del s1, pl prog = Progress(log, sum(not p.exists() for p in paths), f"features {split} shards") # several processes may drain this queue (python features.py starts a helper); keep going until # every shard exists: shards claimed by a live helper are waited for, stale claims are re-taken while True: for i, p in enumerate(paths): if not claim(p): continue try: df = block_rows(split, i) for m in EMB_MODELS: 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, nw) df = pd.concat([df, f], axis=1) df = pd.concat([df, query_context(df, None, stats, s1_country)], axis=1) save_parquet(df, p) finally: release(p) prog.step() missing = [p for p in paths if not p.exists()] if not missing: break log.info(f"features {split}: waiting for {len(missing)} shard(s) held by other workers") time.sleep(30) def feature_columns(df): return [c for c in df.columns if c not in ("q", "c", "y")] # ----------------------------------------------------------------------------------------------- union with ename (v2) CONTEXT_COLS = (["q_ncand", "lex_rel", "q_best_n_tset", "n_tset_gap", "tok_ratio_n", "tok_ratio_a"] + [f"{x}{m}{y}" for m in EMB_MODELS for x, y in (("", "_cos_gap"), ("q_best_", ""), ("", "_cos_rank"), ("", "_cos_qz"), ("", "_cos_cz"), ("q_best_", "_cz"))]) ENAME_K = 5 def run_featx(split): """README v2 §2.2/§3.2: the fused top-K rows plus the ename candidates as a guaranteed union. New (q, c) pairs get exact e5/bge/e5s cosines and the full pair features; within-query context (ranks, gaps, z-scores, counts) is recomputed over the union; flags from_ename / ename_only and ename_s / ename_r are added. Feature shards from run_features are read, never modified.""" from config import CFG, RETR # noqa: F401 from decide import in_sorted from embeddings import load_emb from prep import load_splits from blocking import chosen_retr from retrieve_dense import queries ns = n_shards(split) paths = [shard_path("featx", split, i) for i in range(ns)] if all(p.exists() for p in paths): return retr = chosen_retr() Kc = load_json(WORK / "reports" / "k_choice.json")["K"] Q = queries(split) C = CFG["query_chunk"] stats = {m: load_json(WORK / "reports" / f"cos_stats_{m}_{split}.json") for m in EMB_MODELS} gt = None if split == "train": sp = load_splits() gt = np.sort(sp["gt_q"] * np.int64(1 << 32) + sp["gt_c"]) nw = n_workers() with Pool(nw) as pool: 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} s1_country = s1.country.values del s1, pl embs = {m: (load_emb(m, split, "s1"), load_emb(m, split, "pool")) for m in EMB_MODELS} prog = Progress(log, sum(not p.exists() for p in paths), f"featx {split} shards") stat = {"new_pairs": 0, "ename_pairs": 0, "rows": 0} while True: for i, p in enumerate(paths): if not claim(p): continue try: base = pd.read_parquet(shard_path("feat", split, i)) z = np.load(shard_path("dense", split, i, "npz", "ename")) q = np.repeat(Q[i * C:(i + 1) * C], z["cand"].shape[1]) en = pd.DataFrame({"q": q.astype(np.int32), "c": z["cand"].ravel(), "ename_s": z["score"].astype(np.float32).ravel(), "ename_r": np.tile(np.arange(z["cand"].shape[1]), len(z["cand"])).astype(np.float32)}) en = en[en.c >= 0].reset_index(drop=True) key = lambda d: d.q.values.astype(np.int64) * (1 << 32) + d.c.values bkey = np.sort(key(base)) new = en[~in_sorted(key(en), bkey)][["q", "c"]].reset_index(drop=True) if len(new): new["rrf"] = np.float32(0) new["rrf_rank"] = np.int16(Kc) # sentinel: outside the fused top-K new["support"] = np.int8(0) new["y"] = (in_sorted(key(new), gt).astype(np.int8) if gt is not None else np.int8(-1)) for r in retr: new[f"{r}_r"] = np.int16(CFG["dense_k"]) new[f"{r}_s"] = np.float32(np.nan) new[f"{r}_gap"] = np.float32(np.nan) for m, (es, ep) in embs.items(): a = es.take(new.q.values).astype(np.float32) b = ep.take(new.c.values).astype(np.float32) new[f"{m}_cos"] = (a * b).sum(1) new["tri_max_cos"] = np.float32(np.nan) new["tri_n_close"] = np.float32(0) f = pair_features(pool, s1tab, pooltab, new, nw) new = pd.concat([new, f], axis=1) df = pd.concat([base.drop(columns=[c for c in CONTEXT_COLS if c in base.columns]), new[[c for c in base.columns if c not in CONTEXT_COLS]] if len(new) else None], ignore_index=True) df["ename_only"] = np.r_[np.zeros(len(base), np.int8), np.ones(len(new), np.int8)] df = df.merge(en, on=["q", "c"], how="left") df["from_ename"] = df.ename_s.notna().astype(np.int8) df["ename_r"] = df.ename_r.fillna(ENAME_K).astype(np.float32) df = df.sort_values(["q", "ename_only", "rrf_rank"], kind="stable").reset_index(drop=True) df = pd.concat([df, query_context(df, None, stats, s1_country)], axis=1) save_parquet(df, p) stat["new_pairs"] += len(new); stat["ename_pairs"] += len(en); stat["rows"] += len(df) finally: release(p) prog.step() missing = [p for p in paths if not p.exists()] if not missing: break log.info(f"featx {split}: waiting for {len(missing)} shard(s) held by other workers") time.sleep(30) log.info(f"featx {split} (this process): {stat}") if __name__ == "__main__": # helper: python features.py — joins that stage's shard queue (run_featx if sys.argv[1] == "featx" else run_features)(sys.argv[-1])