| """Stage 1 (cheap LightGBM over all fused candidates), stage-2 candidate set, stage 2 (LightGBM + cross-encoder, |
| calibration, decision selection), test inference + submission files, and the leave-one-country-out proxy.""" |
| import pickle |
| import subprocess |
| import sys |
| from multiprocessing import Pool |
|
|
| import lightgbm as lgb |
| import numpy as np |
| import pandas as pd |
|
|
| from common import get_logger, load_json, save_json, save_parquet, shard_path, n_workers |
| from config import CFG, OUTPUT, PATHS, WORK |
| from decide import crossfit_isotonic, decide, fit_isotonic, macro_f05 |
| from features import feature_columns |
| from prep import load_prep, load_splits |
| from retrieve_dense import n_shards |
|
|
| log = get_logger("stack") |
| MODELS = WORK / "models" |
| PARAMS = dict(objective="binary", learning_rate=0.05, num_leaves=127, min_data_in_leaf=100, feature_fraction=0.8, |
| bagging_fraction=0.8, bagging_freq=1, lambda_l2=1.0, max_bin=255, verbose=-1) |
|
|
|
|
| def fold_of(q, n): |
| return (q.astype(np.int64) * 2654435761 % 2 ** 32) % n |
|
|
|
|
| def train_lgb(X, y, q, path, seed): |
| """Early stopping on an inner holdout of whole S1 groups; model saved atomically for resume.""" |
| if path.exists(): |
| return lgb.Booster(model_file=str(path)) |
| inner = fold_of(q + 7, 10) == 0 |
| p = dict(PARAMS, seed=seed, num_threads=CFG["lgb_threads"]) |
| dtr = lgb.Dataset(X[~inner], y[~inner], free_raw_data=True) |
| dva = lgb.Dataset(X[inner], y[inner], reference=dtr) |
| m = lgb.train(p, dtr, num_boost_round=3000, valid_sets=[dva], |
| callbacks=[lgb.early_stopping(100, verbose=False), lgb.log_evaluation(250)]) |
| path.parent.mkdir(parents=True, exist_ok=True) |
| tmp = path.with_suffix(".tmp") |
| m.save_model(str(tmp), num_iteration=m.best_iteration) |
| tmp.replace(path) |
| log.info(f"saved {path.name} best_iter={m.best_iteration}") |
| return lgb.Booster(model_file=str(path)) |
|
|
|
|
| def load_feats(split, qset=None): |
| parts = [] |
| for i in range(n_shards(split)): |
| d = pd.read_parquet(shard_path("feat", split, i)) |
| if qset is not None: |
| d = d[np.isin(d.q.values, qset)] |
| parts.append(d) |
| return pd.concat(parts, ignore_index=True) |
|
|
|
|
| |
|
|
| def run_stage1_train(): |
| sp = load_splits() |
| df = load_feats("train") |
| feats = feature_columns(df) |
| MODELS.mkdir(parents=True, exist_ok=True) |
| save_json(feats, MODELS / "stage1_features.json") |
| stack = df[np.isin(df.q.values, sp["g_stack"])].reset_index(drop=True) |
| val = df[np.isin(df.q.values, sp["g_val"])].reset_index(drop=True) |
| del df |
| F = CFG["n_folds"] |
| fold = fold_of(stack.q.values, F) |
| oof = np.zeros(len(stack)) |
| vp = np.zeros(len(val)) |
| rng = np.random.default_rng(CFG["seed"]) |
| for k in range(F): |
| tr = stack[fold != k] |
| keep = (tr.y.values == 1) | (tr.rrf_rank.values < CFG["s1_neg_keep_rank"]) | (rng.random(len(tr)) < CFG["s1_neg_frac"]) |
| tr = tr[keep] |
| m = train_lgb(tr[feats].values.astype(np.float32), tr.y.values, tr.q.values, MODELS / f"s1_fold{k}.txt", k) |
| oof[fold == k] = m.predict(stack.loc[fold == k, feats].values.astype(np.float32)) |
| vp += m.predict(val[feats].values.astype(np.float32)) / F |
| out = pd.concat([stack[["q", "c", "y"]].assign(p1=oof), val[["q", "c", "y"]].assign(p1=vp)], ignore_index=True) |
| save_parquet(out, WORK / "stage1" / "train.parquet") |
| from sklearn.metrics import roc_auc_score |
| log.info(f"stage1 AUC stack-oof={roc_auc_score(stack.y, oof):.5f} val={roc_auc_score(val.y, vp):.5f}") |
|
|
|
|
| def fold_predictor(models, X_check): |
| """Mean of the fold models' probabilities. Uses the exact GPU tree evaluator (gpu_lgb.py) when a GPU is |
| present and it reproduces LightGBM on a check sample (|diff| < 1e-9); otherwise LightGBM on the CPU.""" |
| cpu = lambda X: sum(m.predict(X, num_threads=CFG["lgb_threads"]) for m in models) / len(models) |
| try: |
| import torch |
| if torch.cuda.is_available(): |
| from gpu_lgb import GpuForest |
| fs = [GpuForest(m, torch.device("cuda:0")) for m in models] |
| gpu = lambda X: sum(f.predict(X) for f in fs) / len(fs) |
| diff = float(np.abs(gpu(X_check) - cpu(X_check)).max()) |
| if diff < 1e-9: |
| log.info(f"LightGBM scoring on GPU (check max|diff|={diff:.1e})") |
| return gpu |
| log.warning(f"GPU tree evaluator differs from LightGBM ({diff:.1e}); using CPU") |
| except Exception as e: |
| log.warning(f"GPU tree evaluator unavailable ({e}); using CPU") |
| return cpu |
|
|
|
|
| def run_stage1_test(): |
| feats = load_json(MODELS / "stage1_features.json") |
| models = [lgb.Booster(model_file=str(MODELS / f"s1_fold{k}.txt")) for k in range(CFG["n_folds"])] |
| predict = None |
| for i in range(n_shards("test")): |
| out = shard_path("stage1", "test", i) |
| if out.exists(): |
| continue |
| d = pd.read_parquet(shard_path("feat", "test", i), columns=["q", "c"] + feats) |
| X = d[feats].values.astype(np.float32) |
| if predict is None: |
| predict = fold_predictor(models, X[:20_000]) |
| p = predict(X) |
| save_parquet(d[["q", "c"]].assign(p1=p.astype(np.float32)), out) |
|
|
|
|
| |
|
|
| def _s2(df): |
| df = df.sort_values(["q", "p1"], ascending=[True, False], kind="stable") |
| g = df.groupby("q") |
| df["p1_rank"] = g.cumcount().astype(np.float32) |
| df["p1_gap"] = (g.p1.transform("max") - df.p1).astype(np.float32) |
| return df[(df.p1_rank < CFG["s2_topn"]) & (df.p1 >= CFG["s2_min_prob"])].reset_index(drop=True) |
|
|
|
|
| def run_s2set(split): |
| if split == "train": |
| s2 = _s2(pd.read_parquet(WORK / "stage1" / "train.parquet")) |
| sp = load_splits() |
| feat_pos = int(pd.read_parquet(WORK / "stage1" / "train.parquet", columns=["y"]).y.sum()) |
| allq = np.concatenate([sp["g_stack"], sp["g_val"]]) |
| rep = {"rows": len(s2), "true_in_blocking": feat_pos, "true_in_s2": int(s2.y.sum()), |
| "true_total": int(sp["n_true"][allq].sum())} |
| rep["s2_recall_vs_blocking"] = rep["true_in_s2"] / max(feat_pos, 1) |
| rep["s2_recall_vs_total"] = rep["true_in_s2"] / max(rep["true_total"], 1) |
| save_json(rep, WORK / "reports" / "s2set_train.json") |
| log.info(f"stage-2 set: {rep}") |
| else: |
| s2 = _s2(pd.concat([pd.read_parquet(shard_path("stage1", "test", i)) for i in range(n_shards("test"))], |
| ignore_index=True)) |
| save_parquet(s2, WORK / "s2" / f"{split}.parquet") |
|
|
|
|
| |
|
|
| def _attach(split, s2): |
| """Join pair features and cross-encoder scores onto the stage-2 rows.""" |
| from cross_encoder import load_ce_scores |
| qs = np.unique(s2.q.values) |
| parts = [] |
| key = s2[["q", "c"]] |
| for i in range(n_shards(split)): |
| d = pd.read_parquet(shard_path("feat", split, i)) |
| d = d[np.isin(d.q.values, qs)] |
| parts.append(key[np.isin(key.q.values, np.unique(d.q.values))].merge(d, on=["q", "c"], how="left")) |
| f = pd.concat(parts, ignore_index=True) |
| s2 = s2.merge(f.drop(columns=[c for c in ("y",) if c in s2.columns]), on=["q", "c"], how="left") |
| ce = load_ce_scores(split) |
| s2 = s2.merge(ce, on=["q", "c"], how="left") |
| g = s2.groupby("q") |
| s2["ce_rank"] = g.ce.rank(ascending=False, method="first").astype(np.float32) |
| s2["ce_gap"] = (g.ce.transform("max") - s2.ce).astype(np.float32) |
| return s2 |
|
|
|
|
| def stage2_features(s2, variant): |
| base = [c for c in s2.columns if c not in ("q", "c", "y", "ce", "ce_rank", "ce_gap")] |
| return base + (["ce", "ce_rank", "ce_gap"] if variant == "ce" else []) |
|
|
|
|
| def run_stage2_train(): |
| sp = load_splits() |
| s2 = pd.read_parquet(WORK / "s2" / "train.parquet") |
| s2 = _attach("train", s2) |
| stack = s2[np.isin(s2.q.values, sp["g_stack"])].reset_index(drop=True) |
| val = s2[np.isin(s2.q.values, sp["g_val"])].reset_index(drop=True) |
| F = CFG["n_folds"] |
| fold = fold_of(stack.q.values, F) |
| country = load_prep("train", "s1", ["country"]).country.values |
| results, preds = {}, {} |
| with Pool(n_workers()) as pool: |
| for variant in ("noce", "ce"): |
| feats = stage2_features(s2, variant) |
| oof, vp = np.zeros(len(stack)), np.zeros(len(val)) |
| for k in range(F): |
| tr = stack[fold != k] |
| m = train_lgb(tr[feats].values.astype(np.float32), tr.y.values, tr.q.values, |
| MODELS / f"s2_{variant}_fold{k}.txt", 100 + k) |
| oof[fold == k] = m.predict(stack.loc[fold == k, feats].values.astype(np.float32)) |
| vp += m.predict(val[feats].values.astype(np.float32)) / F |
| cal = crossfit_isotonic(oof, stack.y.values, stack.q.values, F) |
| preds[variant] = (oof, vp, feats) |
| d = stack[["q", "c"]].assign(p=cal) |
| fams = [("expf05", None)] + [("threshold", t) for t in CFG["thr_grid"]] |
| for fam, thr in fams: |
| sel = decide(d, fam, thr, CFG["arb_max_iter"], pool) |
| f, _ = macro_f05(sel, sp["g_stack"], sp["n_true"], sp["gt_q"], sp["gt_c"]) |
| results[(variant, fam, thr)] = f |
| log.info(f"stack F0.5 variant={variant} family={fam} thr={thr}: {f:.5f}") |
| best = max(results, key=results.get) |
| variant, fam, thr = best |
| oof, vp, feats = preds[variant] |
| iso = fit_isotonic(oof, stack.y.values) |
| vcal = iso.predict(vp) |
| sel = decide(val[["q", "c"]].assign(p=vcal), fam, thr, CFG["arb_max_iter"], pool) |
| f, per = macro_f05(sel, sp["g_val"], sp["n_true"], sp["gt_q"], sp["gt_c"]) |
| per["country"] = country[per.q.values] |
| by_c = per.groupby("country").f.mean().to_dict() |
| npred = sel.groupby("q").size() |
| val_stats = _pred_stats(sel, sp["g_val"], country) |
| bundle = {"variant": variant, "family": fam, "thr": thr, "features": feats, |
| "stack_scores": {f"{v}|{fa}|{t}": s for (v, fa, t), s in results.items()}, |
| "val_f05": f, "val_f05_by_country": by_c, "val_pred_stats": val_stats} |
| with open(MODELS / "stage2_iso.pkl", "wb") as fh: |
| pickle.dump(iso, fh) |
| save_json(bundle, MODELS / "stage2_bundle.json") |
| |
| imp = pd.DataFrame({"feature": feats}) |
| for k in range(F): |
| m = lgb.Booster(model_file=str(MODELS / f"s2_{variant}_fold{k}.txt")) |
| imp[f"gain_{k}"] = m.feature_importance("gain") |
| imp["gain"] = imp.filter(like="gain_").mean(1) |
| imp.sort_values("gain", ascending=False)[["feature", "gain"]].to_csv(WORK / "reports" / "feature_importance.tsv", sep="\t", index=False) |
| per.sort_values("f").head(2000).to_csv(WORK / "reports" / "val_errors.tsv", sep="\t", index=False) |
| lines = [f"# Stage-2 validation report", "", |
| f"- chosen: variant **{variant}**, decision **{fam}**" + (f" (thr {thr})" if thr else ""), |
| f"- validation macro F0.5: **{f:.5f}**", |
| f"- by country: {', '.join(f'{k}: {v:.5f}' for k, v in by_c.items())}", |
| f"- stack F0.5 of all options: see models/stage2_bundle.json", |
| f"- predicted matches per S1 (val): {val_stats}"] |
| (WORK / "reports" / "report.md").write_text("\n".join(lines) + "\n") |
| log.info("\n".join(lines)) |
|
|
|
|
| def _pred_stats(sel, qs, country): |
| n = pd.Series(0, index=qs).add(sel.groupby("q").size(), fill_value=0).loc[qs] |
| c = pd.Series(country[qs], index=qs) |
| out = {} |
| for cc, g in n.groupby(c): |
| out[cc] = {"n_s1": int(len(g)), "empty_rate": round(float((g == 0).mean()), 4), |
| "mean_pred": round(float(g.mean()), 3), |
| "dist": {str(k): round(float(v), 4) for k, v in g.clip(upper=6).value_counts(normalize=True).sort_index().items()}} |
| return out |
|
|
|
|
| |
|
|
| def run_stage2_test(): |
| bundle = load_json(MODELS / "stage2_bundle.json") |
| with open(MODELS / "stage2_iso.pkl", "rb") as fh: |
| iso = pickle.load(fh) |
| s2 = _attach("test", pd.read_parquet(WORK / "s2" / "test.parquet")) |
| feats = bundle["features"] |
| X = s2[feats].values.astype(np.float32) |
| models = [lgb.Booster(model_file=str(MODELS / f"s2_{bundle['variant']}_fold{k}.txt")) for k in range(CFG["n_folds"])] |
| raw = fold_predictor(models, X[:20_000])(X) |
| s2["p"] = iso.predict(raw) |
| save_parquet(s2[["q", "c", "p1", "p"] + (["ce"] if "ce" in s2 else [])], WORK / "stage2" / "test_scores.parquet") |
| with Pool(n_workers()) as pool: |
| sel = decide(s2[["q", "c", "p"]], bundle["family"], bundle["thr"], CFG["arb_max_iter"], pool) |
| write_submission(s2, sel, bundle) |
|
|
|
|
| def write_submission(s2, sel, bundle): |
| sp = load_splits() |
| s1 = load_prep("test", "s1", ["id", "country"]) |
| pool_ids = load_prep("test", "pool", ["id"]).id.values |
| q_all = sp["q_test"] |
| OUTPUT.mkdir(parents=True, exist_ok=True) |
|
|
| def lists(df): |
| d = df.assign(cid=pool_ids[df.c.values]).groupby("q").cid.apply(lambda x: ",".join(dict.fromkeys(x))) |
| return pd.Series(q_all).map(d).fillna("").values |
|
|
| cand = pd.DataFrame({"source1_entity_id": s1.id.values[q_all], "candidate_entity_ids": lists(s2)}) |
| match = pd.DataFrame({"source1_entity_id": s1.id.values[q_all], "matched_entity_ids": lists(sel)}) |
| for df, name in ((match, "matching_results.tsv"), (cand, "candidate_pairs.tsv")): |
| tmp = OUTPUT / f".{name}.tmp" |
| df.to_csv(tmp, sep="\t", index=False) |
| tmp.replace(OUTPUT / name) |
| |
| assert match.source1_entity_id.is_unique and len(match) == len(q_all) |
| cset = dict(zip(cand.source1_entity_id, cand.candidate_entity_ids)) |
| for sid, m in zip(match.source1_entity_id, match.matched_entity_ids): |
| if m: |
| ids = m.split(",") |
| assert len(ids) == len(set(ids)) and all(x[:3] in ("S2-", "S3-") for x in ids) |
| assert set(ids) <= set(cset[sid].split(",")), f"match outside candidates for {sid}" |
| stats = _pred_stats(sel, q_all, s1.country.values) |
| rep = {"test_pred_stats": stats, "val_pred_stats": bundle["val_pred_stats"], |
| "train_singleton_rate": float((sp["n_true"] == 0).mean()), |
| "train_mean_true": float(sp["n_true"].mean())} |
| val_mean = np.mean([v["mean_pred"] for v in bundle["val_pred_stats"].values()]) |
| val_empty = np.mean([v["empty_rate"] for v in bundle["val_pred_stats"].values()]) |
| flags = [] |
| for c, v in stats.items(): |
| if abs(v["mean_pred"] - val_mean) > 0.35 * val_mean or abs(v["empty_rate"] - val_empty) > 0.05: |
| flags.append(f"{c}: mean_pred={v['mean_pred']} empty_rate={v['empty_rate']} vs val {val_mean:.3f}/{val_empty:.3f}") |
| rep["country_shift_flags"] = flags |
| if CFG["test_limit"] is None: |
| r = subprocess.run([sys.executable, str(PATHS["validator"]), "--matching", str(OUTPUT / "matching_results.tsv"), |
| "--candidate", str(OUTPUT / "candidate_pairs.tsv"), "--test-dir", str(PATHS["raw_test_dir"]), "--check-ids"], |
| capture_output=True, text=True) |
| rep["validator"] = {"returncode": r.returncode, "stdout": r.stdout[-4000:], "stderr": r.stderr[-2000:]} |
| if r.returncode != 0: |
| save_json(rep, WORK / "reports" / "test_report.json") |
| raise RuntimeError(f"validator failed:\n{r.stdout}\n{r.stderr}") |
| save_json(rep, WORK / "reports" / "test_report.json") |
| log.info(f"submission written to {OUTPUT}; per-country prediction stats: {stats}; flags: {flags}") |
|
|
|
|
| |
|
|
| def run_eval_b(): |
| """Leave-one-country-out proxy for unseen France: train on one country's stack, evaluate on the other |
| country's val (stage-1 features, no cross-encoder), compared with same-country val.""" |
| sp = load_splits() |
| country = load_prep("train", "s1", ["country"]).country.values |
| df = load_feats("train") |
| feats = feature_columns(df) |
| rep = {} |
| with Pool(n_workers()) as pool: |
| for tr_c, te_c in (("us", "india"), ("india", "us")): |
| st = df[np.isin(df.q.values, sp["g_stack"]) & (country[df.q.values] == tr_c)].reset_index(drop=True) |
| fold = fold_of(st.q.values, 3) |
| oof = np.zeros(len(st)) |
| ms = [] |
| for k in range(3): |
| tr = st[fold != k] |
| m = train_lgb(tr[feats].values.astype(np.float32), tr.y.values, tr.q.values, |
| MODELS / f"B_{tr_c}_fold{k}.txt", 200 + k) |
| oof[fold == k] = m.predict(st.loc[fold == k, feats].values.astype(np.float32)) |
| ms.append(m) |
| iso = fit_isotonic(oof, st.y.values) |
| for ev_c in (tr_c, te_c): |
| qv = sp["g_val"][country[sp["g_val"]] == ev_c] |
| v = df[np.isin(df.q.values, qv)] |
| p = iso.predict(sum(m.predict(v[feats].values.astype(np.float32)) for m in ms) / 3) |
| sel = decide(v[["q", "c"]].assign(p=p), "expf05", None, CFG["arb_max_iter"], pool) |
| f, _ = macro_f05(sel, qv, sp["n_true"], sp["gt_q"], sp["gt_c"]) |
| rep[f"train_{tr_c}->val_{ev_c}"] = {"f05": f, **_pred_stats(sel, qv, country)[ev_c]} |
| log.info(f"split B train {tr_c} -> val {ev_c}: F0.5={f:.5f}") |
| save_json(rep, WORK / "reports" / "split_B.json") |
|
|
