| """Test-set inference -> output/matching_results.tsv + output/candidate_pairs.tsv. |
| |
| python src/run_test.py --gpu 1 --artifacts work/results_A |
| |
| Requires (produced earlier in the queue): |
| * work/emb/<model>/test_pool/ and work/emb/<model>/test_s1/ (src/embed.py) |
| * an artifacts dir from run_train_eval.py holding abbrev_map.json, the LightGBM |
| model, the isotonic calibrator, the tuned decision parameter and miss_mass. |
| |
| Output contract (see the organisers' README + utils/validate_submission.py): |
| * exactly one row per test Source-1 entity, INCLUDING singletons (empty list) |
| * ids may only be S2-/S3- ids that exist in the test set |
| * no duplicate ids inside a list, no duplicate source1_entity_id rows |
| * matching_results must be a SUBSET of candidate_pairs |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| import time |
|
|
| import numpy as np |
| import pandas as pd |
|
|
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
|
|
| import blocking as B |
| import data as D |
| import decide as DC |
| import text_utils as T |
| from config import CFG, EMB, PATHS, emb_dir |
| from encoder import load_matrix, read_ids |
| from features import compute_features |
|
|
| T0 = time.time() |
|
|
|
|
| def log(m): |
| print(f"[{time.time() - T0:7.0f}s] {m}", flush=True) |
|
|
|
|
| def write_tsv(path, rows, col): |
| """rows: list of (s1_id, [ids]). Writes the exact organiser format.""" |
| os.makedirs(os.path.dirname(path), exist_ok=True) |
| tmp = path + ".tmp" |
| with open(tmp, "w") as f: |
| f.write(f"source1_entity_id\t{col}\n") |
| for s1, ids in rows: |
| f.write(f"{s1}\t{','.join(ids)}\n") |
| os.replace(tmp, path) |
| log(f"wrote {path} ({len(rows):,} rows)") |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--gpu", default=os.environ.get("M3_GPU")) |
| ap.add_argument("--artifacts", required=True, help="dir written by run_train_eval.py") |
| ap.add_argument("--out", default=PATHS["output_dir"]) |
| ap.add_argument("--chunk", type=int, default=200_000, help="test S1 rows per batch") |
| a = ap.parse_args() |
| if a.gpu is None: |
| sys.exit("pass --gpu 0 / --gpu 1 / --gpu cpu (or set M3_GPU)") |
| device = "cpu" if str(a.gpu) == "cpu" else f"cuda:{int(a.gpu)}" |
| art = a.artifacts |
|
|
| |
| import joblib |
| amap = json.load(open(os.path.join(art, "abbrev_map.json"))) |
| bundle = joblib.load(os.path.join(art, "model_bundle.joblib")) |
| lgbm, iso = bundle["lgbm"], bundle["calibrator"] |
| feat_cols, method = bundle["feat_cols"], bundle["method"] |
| param, miss_mass = bundle["param"], bundle["miss_mass"] |
| country_filter = bundle["country_filter"] |
| use_ce = bundle.get("use_cross_encoder", False) |
| ce_dir = bundle.get("ce_dir") |
| log(f"artifacts: {len(amap)} abbrev rules | {len(feat_cols)} features | " |
| f"{method}({param}) | country_filter={country_filter} | ce={use_ce}") |
|
|
| |
| log("loading test data") |
| s1, s23, _ = D.load_data(PATHS["test"], with_gt=False) |
| log(f"test S1 {len(s1):,} test S2+S3 {len(s23):,}") |
| s1["name_core_abbr"] = [T.apply_abbrev(x, amap) for x in s1["name_core"].values] |
| s23["name_core_abbr"] = [T.apply_abbrev(x, amap) for x in s23["name_core"].values] |
|
|
| def vals(df): |
| v = {c: df[c].values.astype(object) for c in |
| ["entity_id", "name_core_abbr", "name_norm", "address_abbrev_norm", |
| "legal_suffix", "country_norm"]} |
| nums = df["address_numbers"].values.astype(object) |
| raw = df["business_address"].values.astype(object) |
| v["nums"] = np.where(nums == "", raw, nums) |
| return v |
| s1v, s23v = vals(s1), vals(s23) |
|
|
| |
| pool_dir, q_dir = emb_dir("test_pool"), emb_dir("test_s1") |
| c_emb, cmeta = load_matrix(pool_dir) |
| q_emb, qmeta = load_matrix(q_dir) |
| pool_ids, q_ids = read_ids(os.path.join(pool_dir, "ids.txt")), read_ids(os.path.join(q_dir, "ids.txt")) |
| |
| assert list(s23["entity_id"]) == pool_ids, "test pool embedding row order != S2+S3 order" |
| assert list(s1["entity_id"]) == q_ids, "test S1 embedding row order != S1 order" |
| log(f"embeddings ok: pool {c_emb.shape} queries {q_emb.shape} (model {cmeta['model']})") |
|
|
| |
| log("building lexical index over the test pool") |
| lidx = B.LexicalIndex([B.lexical_tokens(n, ad) for n, ad in |
| zip(s23v["name_core_abbr"], s23v["address_abbrev_norm"])], CFG) |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| emb2_on = bool(bundle.get("emb2_enabled", CFG.get("emb2_enabled"))) |
| if emb2_on: |
| import emb_source |
| c2_emb, q2_emb = emb_source.load_emb2("test", len(s1), len(s23)) |
| log(f"second dense channel: pool {c2_emb.shape} (e5-large, pre-computed)") |
|
|
| |
| cascade_m = int(CFG.get("cascade_m", 0) or 0) |
| preranker = None |
| if cascade_m: |
| pr_path = os.path.join(art, "preranker.joblib") |
| if os.path.exists(pr_path): |
| preranker = joblib.load(pr_path) |
| log(f"cascade ON: reranking top-{cascade_m}/entity using a " |
| f"{len(preranker['cols'])}-feature pre-ranker " |
| f"(validated delta F0.5 +0.000000)") |
| else: |
| log(f"cascade requested (M={cascade_m}) but {pr_path} is missing - " |
| f"reranking ALL candidates instead") |
|
|
| ck_dir = os.path.join(art, "test_chunks") |
| os.makedirs(ck_dir, exist_ok=True) |
| n = len(s1) |
| chunk_los_all = list(range(0, n, a.chunk)) |
| _lo = int(os.environ.get("M3_CHUNK_LO", "0")) |
| _hi = int(os.environ.get("M3_CHUNK_HI", str(len(chunk_los_all)))) |
| chunk_los = chunk_los_all[_lo:_hi] |
| skip_decide = os.environ.get("M3_SKIP_DECIDE", "0") == "1" |
| if (_lo, _hi) != (0, len(chunk_los_all)) or skip_decide: |
| log(f"worker slice: chunks [{_lo}:{_hi}] of {len(chunk_los_all)}" |
| f"{' (skipping global decision)' if skip_decide else ''}") |
| for lo in chunk_los: |
| hi = min(lo + a.chunk, n) |
| cpath = os.path.join(ck_dir, f"chunk_{lo:09d}.parquet") |
| if os.path.exists(cpath): |
| log(f"chunk {lo:,}-{hi:,}: already complete, skipping") |
| continue |
| q_rows = np.arange(lo, hi) |
| log(f"chunk {lo:,}-{hi:,}: dense search") |
| emb_s, emb_i = B.dense_topk(np.asarray(q_emb[lo:hi]), c_emb, CFG, device) |
| log(f"chunk {lo:,}-{hi:,}: lexical search") |
| lex = B.lexical_topk(lidx, [B.lexical_tokens(s1v["name_core_abbr"][i], |
| s1v["address_abbrev_norm"][i]) |
| for i in q_rows], CFG["lex_n_jobs"]) |
| emb2_s = emb2_i = None |
| if emb2_on: |
| cfg2 = dict(CFG); cfg2["emb_k"] = CFG["emb2_k"] |
| emb2_s, emb2_i = B.dense_topk(np.asarray(q2_emb[lo:hi]), c2_emb, cfg2, device) |
| cand = B.build_candidates(q_rows, emb_s, emb_i, lex, s1v["country_norm"], |
| s23v["country_norm"], country_filter, |
| emb2_s=emb2_s, emb2_i=emb2_i) |
| qpos = {r: i for i, r in enumerate(q_rows.tolist())} |
| cand = B.fill_missing_emb(cand, lambda rows: q_emb[[lo + qpos[r] for r in rows]], c_emb) |
| log(f"chunk {lo:,}-{hi:,}: {len(cand):,} candidates; features") |
| feats = compute_features(cand, s1v, s23v, CFG["feat_n_jobs"]) |
| cand = pd.concat([cand, feats], axis=1) |
| del feats |
|
|
| if preranker is not None: |
| before = len(cand) |
| sc = preranker["model"].predict(cand[preranker["cols"]]) |
| rk = pd.Series(sc, index=cand.index).groupby(cand["q"].values) \ |
| .rank(ascending=False, method="first") |
| cand = cand[rk.values <= cascade_m].reset_index(drop=True) |
| log(f"chunk {lo:,}-{hi:,}: cascade kept {len(cand):,}/{before:,} " |
| f"({100*len(cand)/max(before,1):.1f}%) for reranking") |
|
|
| if use_ce: |
| from cross_encoder import pair_text, predict_cross_encoder |
| log(f"chunk {lo:,}-{hi:,}: cross-encoder") |
| ta = [pair_text(s1["business_name"].values[r], s1["business_address"].values[r], |
| s1["country"].values[r]) for r in cand["q"].values] |
| tb = [pair_text(s23["business_name"].values[c], s23["business_address"].values[c], |
| s23["country"].values[c]) for c in cand["c"].values] |
| |
| cand["ce_prob"] = predict_cross_encoder( |
| ta, tb, ce_dir, CFG, cache=os.path.join(ck_dir, f"ce_{lo:09d}")) |
| del ta, tb |
|
|
| cand["p"] = iso.predict(lgbm.predict(cand[feat_cols])) |
| out = pd.DataFrame({"q": cand["q"].values.astype(np.int64), |
| "c": cand["c"].values.astype(np.int64), |
| "p": cand["p"].values.astype(np.float32)}) |
| out.to_parquet(cpath + ".tmp", index=False) |
| os.replace(cpath + ".tmp", cpath) |
| for stale in (f"ce_{lo:09d}.part.npy", f"ce_{lo:09d}.done.json"): |
| sp = os.path.join(ck_dir, stale) |
| if os.path.exists(sp): |
| os.remove(sp) |
| log(f"chunk {lo:,}-{hi:,}: persisted {len(out):,} scored pairs -> {cpath}") |
| del cand, out |
|
|
| |
| if skip_decide: |
| log(f"worker finished its {len(chunk_los)} chunk(s); global decision skipped " |
| f"(run again without M3_SKIP_DECIDE once every chunk exists)") |
| return |
|
|
| |
| |
| |
| missing = [lo for lo in chunk_los_all |
| if not os.path.exists(os.path.join(ck_dir, f"chunk_{lo:09d}.parquet"))] |
| if missing: |
| raise SystemExit(f"cannot run the global decision: {len(missing)} chunk(s) " |
| f"still missing (first at row {missing[0]:,}). Wait for the " |
| f"other worker(s) to finish.") |
| log("loading every chunk for GLOBAL conflict resolution + selection") |
| allc = pd.concat([pd.read_parquet(os.path.join(ck_dir, f"chunk_{lo:09d}.parquet")) |
| for lo in chunk_los_all], ignore_index=True) |
| log(f"{len(allc):,} scored pairs over {n:,} entities " |
| f"({len(allc)/max(n,1):.1f} candidates/entity)") |
| n_claims = len(allc) |
| n_owners = allc["c"].nunique() |
| log(f"{n_owners:,} distinct pool records claimed by {n_claims:,} candidate edges") |
| preds = DC.predict_sets(allc, "q", "c", "p", method, param, miss_mass, conflicts=True) |
| log(f"selection done: {len(preds):,} entities received at least one match") |
|
|
| pool_id_arr = s23v["entity_id"] |
| s1_id_arr = s1v["entity_id"] |
| by_q = allc.groupby("q")["c"].apply(list).to_dict() |
| del allc |
|
|
| match_rows, cand_rows = [], [] |
| for r in range(n): |
| sid = s1_id_arr[r] |
| cand_rows.append((sid, [str(pool_id_arr[c]) for c in by_q.get(r, [])])) |
| match_rows.append((sid, [str(pool_id_arr[c]) for c in preds.get(r, [])])) |
|
|
| |
| assert len(match_rows) == n, f"{len(match_rows)} rows for {n} test S1 entities" |
| valid = set(s23v["entity_id"].astype(str)) |
| seen_owner = {} |
| for sid, ids in match_rows: |
| assert len(ids) == len(set(ids)), f"duplicate ids for {sid}" |
| assert all(i in valid for i in ids), f"unknown id in {sid}" |
| for i in ids: |
| assert i not in seen_owner, f"{i} claimed by both {seen_owner[i]} and {sid}" |
| seen_owner[i] = sid |
| cand_map = {s: set(i) for s, i in cand_rows} |
| for sid, ids in match_rows: |
| assert set(ids) <= cand_map[sid], f"{sid}: match not a subset of candidates" |
| log(f"verified: {len(seen_owner):,} matched ids, each owned by exactly one S1 entity") |
|
|
| write_tsv(os.path.join(a.out, "matching_results.tsv"), match_rows, "matched_entity_ids") |
| write_tsv(os.path.join(a.out, "candidate_pairs.tsv"), cand_rows, "candidate_entity_ids") |
| nz = sum(1 for _, i in match_rows if i) |
| log(f"predicted matches for {nz:,}/{n:,} entities " |
| f"({1 - nz / n:.4f} predicted-singleton rate; train singleton rate is 0.0558)") |
| log(f"now run: python3 {PATHS['validator']} --matching {a.out}/matching_results.tsv " |
| f"--candidate {a.out}/candidate_pairs.tsv --test-dir {os.path.dirname(PATHS['test']['source1'])}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
