model-4/Colab/src/features.py

"""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")]