#!/usr/bin/env python3 """ Option B — Production Entity Resolution Pipeline Amazon ML Challenge 2026 Single-file, end-to-end pipeline: 1. Transliteration & normalization (Indic scripts + French diacritics) 2. DuckDB multi-view blocking (S2+S3 combined, indexed once per country) 3. LightGBM classifier with 25 rich lexical/structural features 4. Per-country threshold optimization on OOF predictions 5. 1:Many target-exclusivity resolution (each target → 1 S1; each S1 → many targets) 6. Full test-set inference on all 1.73M test entities 7. Output + validation Run on Kaggle with 2x T4 GPUs: !python run_option_b_final.py """ import os, sys, gc, re, time, json, unicodedata os.environ['PYTHONUNBUFFERED'] = '1' if hasattr(sys.stdout, 'reconfigure'): sys.stdout.reconfigure(line_buffering=True) from pathlib import Path from collections import defaultdict import numpy as np import pandas as pd import duckdb from rapidfuzz import fuzz # ───────────────────────────────────────────────────────── # 0. CONFIG # ───────────────────────────────────────────────────────── SAMPLE_N_PER_COUNTRY = 100_000 # training sample size per country TOP_K = 50 # max candidates per S1 entity (raised from 35 for recall) FREQ_CAP_DEFAULT = 800 # raised from 500 for better recall on common tokens FREQ_CAP_FRANCE = 1500 S1_BATCH_SIZE = 50_000 # S1 entities per DuckDB blocking batch # Paths — auto-detect Kaggle vs local DATA_DIR = None for p in [Path("/kaggle/input/datasets/arghasaha1905/amazon-ml"), Path("/kaggle/input/amazon-ml-challenge-2026"), Path("/kaggle/input/amazon-ml"), Path("dataset")]: if p.exists(): DATA_DIR = p break assert DATA_DIR is not None, "Cannot find dataset directory" OUT_DIR = Path("/kaggle/working/output") if Path("/kaggle/working").exists() else Path("output") CACHE_DIR = Path("/kaggle/working/cache") if Path("/kaggle/working").exists() else Path("cache") OUT_DIR.mkdir(parents=True, exist_ok=True) CACHE_DIR.mkdir(parents=True, exist_ok=True) print(f"DATA_DIR = {DATA_DIR}") print(f"OUT_DIR = {OUT_DIR}") print(f"CACHE_DIR = {CACHE_DIR}") # ───────────────────────────────────────────────────────── # 1. TRANSLITERATION # ───────────────────────────────────────────────────────── DEV_PHONETICS = { 0x05:'a',0x06:'a',0x07:'i',0x08:'i',0x09:'u',0x0A:'u',0x0B:'ri', 0x0E:'e',0x0F:'e',0x10:'ai',0x12:'o',0x13:'o',0x14:'au', 0x15:'k',0x16:'kh',0x17:'g',0x18:'gh',0x19:'n', 0x1A:'ch',0x1B:'chh',0x1C:'j',0x1D:'jh',0x1E:'n', 0x1F:'t',0x20:'th',0x21:'d',0x22:'dh',0x23:'n', 0x24:'t',0x25:'th',0x26:'d',0x27:'dh',0x28:'n', 0x2A:'p',0x2B:'ph',0x2C:'b',0x2D:'bh',0x2E:'m', 0x2F:'y',0x30:'r',0x31:'r',0x32:'l',0x33:'l',0x35:'v', 0x36:'sh',0x37:'sh',0x38:'s',0x39:'h', 0x3E:'a',0x3F:'i',0x40:'i',0x41:'u',0x42:'u',0x43:'ri', 0x46:'e',0x47:'e',0x48:'ai',0x4A:'o',0x4B:'o',0x4C:'au', 0x4D:'',0x02:'n',0x01:'n',0x03:'h' } TRANSLIT_MAP = {chr(base+off): lat for base in (0x0900,0x0980,0x0A00,0x0A80,0x0B00, 0x0B80,0x0C00,0x0C80,0x0D00) for off, lat in DEV_PHONETICS.items()} TRANSLIT_MAP.update({ 'é':'e','è':'e','ê':'e','ë':'e','à':'a','â':'a','ä':'a', 'ô':'o','ö':'o','î':'i','ï':'i','ù':'u','û':'u','ü':'u', 'ç':'c','ñ':'n', 'É':'e','È':'e','Ê':'e','Ë':'e','À':'a','Â':'a','Ä':'a', 'Ô':'o','Ö':'o','Î':'i','Ï':'i','Ù':'u','Û':'u','Ü':'u', 'Ç':'c','Ñ':'n', '«':'','»':'','\u2018':"'",'\u201c':'"','\u201d':'"' }) TRANSLIT_TABLE = str.maketrans(TRANSLIT_MAP) SUFFIX_RE = re.compile( r'\b(pvt|ltd|limited|private|inc|incorporated|llc|corp|corporation|' r'co|company|enterprises|services|solutions|group|sarl|sas|sasu|' r'eurl|sci|association|fils)\b', re.I) _NONALNUM = re.compile(r'[^a-z0-9]+') def fast_norm(value): if not value or (isinstance(value, float) and value != value): return '' s = str(value).translate(TRANSLIT_TABLE) s = unicodedata.normalize('NFKD', s) s = ''.join(ch for ch in s if not unicodedata.combining(ch)) s = s.encode('ascii', 'ignore').decode('ascii').lower() s = re.sub(r'\b(https?://|www\.)', ' ', s) return _NONALNUM.sub(' ', s).strip() def root_name(n): return SUFFIX_RE.sub(' ', n).strip() def extract_postal(addr): m = re.search(r'\b\d{5,6}\b', addr or '') return m.group(0) if m else '' def extract_house(addr): m = re.search(r'\b\d+[a-z]?\b', addr or '') return m.group(0) if m else '' # Sanity check for raw in ["रियल कंसल्टेंट्स प्राइवेट लिमिटेड", "Thérapy", "பாலாஜி"]: out = fast_norm(raw) assert out != '', f"Transliteration FAILED for {raw!r}" print("Transliteration sanity check: PASSED") # ───────────────────────────────────────────────────────── # 2. DATA LOADING # ───────────────────────────────────────────────────────── def find_file(name): for sub in ['', 'train', 'test']: p = DATA_DIR / sub / name if sub else DATA_DIR / name if p.exists(): return p raise FileNotFoundError(f"Cannot find {name} under {DATA_DIR}") def load_tsv(name): path = find_file(name) df = pd.read_csv(path, sep='\t', dtype=str, keep_default_na=False) print(f" {name}: {len(df):,} rows") return df print("\n=== Loading datasets ===") t0_global = time.time() train_s1 = load_tsv("train_source1.tsv") train_s2 = load_tsv("train_source2.tsv") train_s3 = load_tsv("train_source3.tsv") train_gt = load_tsv("train_ground_truth.tsv") # ───────────────────────────────────────────────────────── # 3. SAMPLE + NORMALIZE # ───────────────────────────────────────────────────────── print(f"\n=== Sampling {SAMPLE_N_PER_COUNTRY:,} per country for training ===") train_s1_sample = (train_s1 .groupby('country', group_keys=False) .apply(lambda g: g.sample(min(SAMPLE_N_PER_COUNTRY, len(g)), random_state=42)) .reset_index(drop=True)) keep_ids = set(train_s1_sample['entity_id']) train_gt_sample = train_gt[train_gt['source1_entity_id'].isin(keep_ids)].reset_index(drop=True) print(f"Sampled train_s1: {len(train_s1_sample):,} rows") print(f"Sampled train_gt: {len(train_gt_sample):,} rows") def normalize_df(df, label=''): t = time.time() df = df.copy() df['name_clean'] = [fast_norm(v) for v in df['business_name']] df['name_root'] = [root_name(v) for v in df['name_clean']] df['addr_clean'] = [fast_norm(v) for v in df['business_address']] df['name_compact'] = [_NONALNUM.sub('', v) for v in df['name_clean']] df['house_no'] = [extract_house(v) for v in df['addr_clean']] df['postal'] = [extract_postal(v) for v in df['addr_clean']] print(f" Normalized {label}: {len(df):,} rows in {time.time()-t:.1f}s") return df print("\nNormalizing train S1 sample...") train_s1_sample = normalize_df(train_s1_sample, 'train_s1_sample') empty_frac = (train_s1_sample['name_clean'] == '').mean() print(f"Empty name_clean fraction: {empty_frac:.4%}") assert empty_frac < 0.02, f"Too many empty names: {empty_frac:.2%}" # ───────────────────────────────────────────────────────── # 4. DuckDB BLOCKING # ───────────────────────────────────────────────────────── KEY_SQL_TEMPLATE = """ SELECT entity_id, 'p5:' || substr(name_compact, 1, 5) AS bkey FROM {t} WHERE length(name_compact) >= 5 UNION ALL SELECT entity_id, 'nt:' || tok AS bkey FROM (SELECT entity_id, unnest(string_split(name_clean, ' ')) AS tok FROM {t}) WHERE length(tok) >= 3 UNION ALL SELECT entity_id, 'nb:' || toks[1] || '_' || toks[2] AS bkey FROM (SELECT entity_id, string_split(name_clean, ' ') AS toks FROM {t}) WHERE len(toks) >= 2 AND length(toks[1]) >= 3 AND length(toks[2]) >= 3 UNION ALL SELECT entity_id, 'np:' || substr(toks[1],1,3) || '_' || substr(toks[2],1,3) AS bkey FROM (SELECT entity_id, string_split(name_clean, ' ') AS toks FROM {t}) WHERE len(toks) >= 2 AND length(toks[1]) >= 3 AND length(toks[2]) >= 3 UNION ALL SELECT entity_id, 'at:' || tok AS bkey FROM (SELECT entity_id, unnest(string_split(addr_clean, ' ')) AS tok FROM {t}) WHERE length(tok) >= 4 UNION ALL SELECT entity_id, 'pn:' || postal || '_' || substr(name_compact,1,3) AS bkey FROM {t} WHERE postal != '' AND length(name_compact) >= 3 UNION ALL SELECT entity_id, 'hn:' || house_no || '_' || substr(name_compact,1,3) AS bkey FROM {t} WHERE house_no != '' AND length(name_compact) >= 3 UNION ALL SELECT entity_id, 'ht:' || house_no || '_' || toks[1] AS bkey FROM (SELECT entity_id, house_no, string_split(name_clean, ' ') AS toks FROM {t}) WHERE house_no != '' AND len(toks) >= 1 AND length(toks[1]) >= 3 """ def block_country(con, s1_df, targets_df, country, freq_cap, top_k): """Run DuckDB blocking for one country. Target keys indexed ONCE.""" t_start = time.time() n_s1 = len(s1_df) print(f"\n [{country}] S1={n_s1:,}, Targets={len(targets_df):,}") con.register('tgt_raw', targets_df) con.execute(f"CREATE OR REPLACE TEMP TABLE target_keys AS {KEY_SQL_TEMPLATE.format(t='tgt_raw')}") con.execute(f""" CREATE OR REPLACE TEMP TABLE filtered_tgt_keys AS SELECT t.entity_id, t.bkey FROM target_keys t JOIN (SELECT bkey, count(*) AS n FROM target_keys GROUP BY bkey HAVING n <= {freq_cap}) f ON t.bkey = f.bkey """) con.execute("CREATE INDEX IF NOT EXISTS ftk_idx ON filtered_tgt_keys(bkey);") con.unregister('tgt_raw') entity_ids = s1_df['entity_id'].tolist() n_batches = (n_s1 + S1_BATCH_SIZE - 1) // S1_BATCH_SIZE all_cands = [] for b in range(n_batches): bt = time.time() batch = s1_df.iloc[b * S1_BATCH_SIZE : (b + 1) * S1_BATCH_SIZE] con.register('s1_b', batch) con.execute(f"CREATE OR REPLACE TEMP TABLE s1_keys AS {KEY_SQL_TEMPLATE.format(t='s1_b')}") joined = con.execute(f""" WITH pair_hits AS ( SELECT s.entity_id AS s1_id, t.entity_id AS target_id, count(*) AS key_hits FROM s1_keys s JOIN filtered_tgt_keys t ON s.bkey = t.bkey GROUP BY 1, 2 ), ranked AS ( SELECT *, row_number() OVER(PARTITION BY s1_id ORDER BY key_hits DESC, target_id) AS rk FROM pair_hits ) SELECT s1_id, target_id, key_hits FROM ranked WHERE rk <= {top_k}; """).fetchdf() con.unregister('s1_b') all_cands.append(joined) print(f" Batch {b+1}/{n_batches}: {len(batch):,} S1 -> {len(joined):,} pairs ({time.time()-bt:.1f}s)") con.execute("DROP TABLE IF EXISTS target_keys; DROP TABLE IF EXISTS filtered_tgt_keys; DROP TABLE IF EXISTS s1_keys;") result = pd.concat(all_cands, ignore_index=True) if all_cands else pd.DataFrame(columns=['s1_id','target_id','key_hits']) print(f" [{country}] done: {len(result):,} pairs in {time.time()-t_start:.1f}s") return result def run_full_blocking(s1_df, s2_df, s3_df, label="train"): """Run blocking across all countries, combining S2+S3 targets.""" print(f"\n{'='*60}") print(f"BLOCKING ({label.upper()})") print(f"{'='*60}") t_block = time.time() countries = sorted(s1_df['country'].unique()) all_frames = [] for ctry in countries: con = duckdb.connect() con.execute("PRAGMA threads=8;") con.execute("PRAGMA max_memory='12GB';") cap = FREQ_CAP_FRANCE if ctry == 'France' else FREQ_CAP_DEFAULT s1_c = s1_df[s1_df['country'] == ctry] # Normalize S2+S3 for this country and combine s2_c = normalize_df(s2_df[s2_df['country'] == ctry], f'{label}_s2_{ctry}') s3_c = normalize_df(s3_df[s3_df['country'] == ctry], f'{label}_s3_{ctry}') targets_c = pd.concat([s2_c, s3_c], ignore_index=True) cands = block_country(con, s1_c, targets_c, ctry, cap, TOP_K) all_frames.append(cands) con.close() del s2_c, s3_c, targets_c, con gc.collect() result = pd.concat(all_frames, ignore_index=True) print(f"\nTotal {label} candidate pairs: {len(result):,} in {time.time()-t_block:.1f}s") return result # ───────────────────────────────────────────────────────── # 5. BLOCKING RECALL CEILING CHECK # ───────────────────────────────────────────────────────── def check_recall(candidates, gt_df): gt_pairs = set() for _, row in gt_df.iterrows(): s1 = row['source1_entity_id'] matched = str(row.get('matched_entity_ids', '') or '').strip() if matched and matched != 'nan': for t in matched.split(','): t = t.strip() if t: gt_pairs.add((s1, t)) cand_set = set(zip(candidates['s1_id'], candidates['target_id'])) recalled = gt_pairs & cand_set recall = len(recalled) / len(gt_pairs) if gt_pairs else 0.0 print(f"\n Ground truth pairs: {len(gt_pairs):,}") print(f" Recalled by blocking: {len(recalled):,}") print(f" BLOCKING RECALL CEILING: {recall:.4%}") return recall # ───────────────────────────────────────────────────────── # 6. FEATURE EXTRACTION # ───────────────────────────────────────────────────────── FEATURE_NAMES = [ 'key_hits', 'name_ratio', 'name_tsort', 'name_tset', 'name_partial', 'root_ratio', 'root_tsort', 'addr_ratio', 'addr_tsort', 'addr_tset', 'addr_partial', 'postal_match', 'postal_conflict', 'house_match', 'name_len_ratio', 'addr_token_overlap', ] def extract_features(candidates, s1_df, s2_df, s3_df): """Extract 16 lexical/structural features for candidate pairs.""" t_feat = time.time() print(f"\n Extracting features for {len(candidates):,} pairs...") # Build lookup dicts s1_dict = s1_df.set_index('entity_id')[['name_clean','name_root','addr_clean','house_no','postal','country']].to_dict('index') tgt_df = pd.concat([s2_df, s3_df], ignore_index=True).drop_duplicates(subset='entity_id') tgt_dict = tgt_df.set_index('entity_id')[['name_clean','name_root','addr_clean','house_no','postal','country']].to_dict('index') n = len(candidates) feat = np.zeros((n, len(FEATURE_NAMES)), dtype=np.float32) s1_ids = candidates['s1_id'].values tgt_ids = candidates['target_id'].values hits = candidates['key_hits'].values.astype(np.float32) for i in range(n): s = s1_dict.get(s1_ids[i]) t = tgt_dict.get(tgt_ids[i]) if s is None or t is None: feat[i, 0] = hits[i] continue sn, tn = s['name_clean'], t['name_clean'] sr, tr = s['name_root'], t['name_root'] sa, ta = s['addr_clean'], t['addr_clean'] sp, tp = s['postal'], t['postal'] sh, th = s['house_no'], t['house_no'] feat[i, 0] = hits[i] feat[i, 1] = fuzz.ratio(sn, tn) feat[i, 2] = fuzz.token_sort_ratio(sn, tn) feat[i, 3] = fuzz.token_set_ratio(sn, tn) feat[i, 4] = fuzz.partial_ratio(sn, tn) feat[i, 5] = fuzz.ratio(sr, tr) feat[i, 6] = fuzz.token_sort_ratio(sr, tr) feat[i, 7] = fuzz.ratio(sa, ta) feat[i, 8] = fuzz.token_sort_ratio(sa, ta) feat[i, 9] = fuzz.token_set_ratio(sa, ta) feat[i, 10] = fuzz.partial_ratio(sa, ta) feat[i, 11] = 1.0 if (sp and tp and sp == tp) else 0.0 feat[i, 12] = 1.0 if (sp and tp and sp != tp) else 0.0 feat[i, 13] = 1.0 if (sh and th and sh == th) else 0.0 ln_s, ln_t = max(len(sn), 1), max(len(tn), 1) feat[i, 14] = min(ln_s, ln_t) / max(ln_s, ln_t) s_toks = set(sa.split()) if sa else set() t_toks = set(ta.split()) if ta else set() denom = max(len(s_toks | t_toks), 1) feat[i, 15] = len(s_toks & t_toks) / denom if (i + 1) % 500_000 == 0: print(f" {i+1:,}/{n:,} features extracted ({time.time()-t_feat:.0f}s)") print(f" Features done: {n:,} pairs in {time.time()-t_feat:.1f}s") return feat # ───────────────────────────────────────────────────────── # 7. BUILD GROUND-TRUTH LOOKUP # ───────────────────────────────────────────────────────── gt_lookup = defaultdict(set) for _, row in train_gt_sample.iterrows(): s1 = row['source1_entity_id'] matched = str(row.get('matched_entity_ids', '') or '').strip() if matched and matched != 'nan': for t in matched.split(','): t = t.strip() if t: gt_lookup[s1].add(t) # ───────────────────────────────────────────────────────── # 8. METRIC # ───────────────────────────────────────────────────────── def macro_f05(eval_s1_set, s1_arr, tgt_arr, probs, threshold, gt_map): """Compute Macro F0.5 strictly on eval_s1_set entities.""" preds = defaultdict(set) for s, t, p in zip(s1_arr, tgt_arr, probs): if p >= threshold: preds[s].add(t) scores = [] for s in eval_s1_set: pred = preds.get(s, set()) true = gt_map.get(s, set()) if not pred and not true: scores.append(1.0) continue tp = len(pred & true) prec = tp / len(pred) if pred else 0.0 rec = tp / len(true) if true else 0.0 denom = 0.25 * prec + rec f05 = (1.25 * prec * rec) / denom if denom > 0 else 0.0 scores.append(f05) return float(np.mean(scores)) # ───────────────────────────────────────────────────────── # MAIN PIPELINE # ───────────────────────────────────────────────────────── def main(): global train_s1_sample # need to access for country mapping later # ── PHASE 1: Training Blocking ── train_cands = run_full_blocking(train_s1_sample, train_s2, train_s3, label="train") recall_ceil = check_recall(train_cands, train_gt_sample) if recall_ceil < 0.70: print(f"WARNING: Blocking recall {recall_ceil:.2%} is low. Continuing anyway — the model can still learn.") elif recall_ceil < 0.80: print(f"NOTE: Blocking recall {recall_ceil:.2%} is below 80% but above 70%. Proceeding.") else: print(f"GOOD: Blocking recall {recall_ceil:.2%} is healthy.") # ── PHASE 2: Training Features ── print(f"\n{'='*60}") print("FEATURE EXTRACTION (TRAIN)") print(f"{'='*60}") # Normalize targets that appear in candidate pairs needed_tids = set(train_cands['target_id']) s2_needed = normalize_df(train_s2[train_s2['entity_id'].isin(needed_tids)], 'train_s2_needed') s3_needed = normalize_df(train_s3[train_s3['entity_id'].isin(needed_tids)], 'train_s3_needed') X_train = extract_features(train_cands, train_s1_sample, s2_needed, s3_needed) s1_arr = train_cands['s1_id'].values tgt_arr = train_cands['target_id'].values y_train = np.array([1 if t in gt_lookup[s] else 0 for s, t in zip(s1_arr, tgt_arr)], dtype=np.float32) n_pos = int(y_train.sum()) n_neg = len(y_train) - n_pos scale_pos_weight = n_neg / max(n_pos, 1) print(f"\n Positives: {n_pos:,} | Negatives: {n_neg:,} | scale_pos_weight: {scale_pos_weight:.2f}") del s2_needed, s3_needed gc.collect() # ── PHASE 3: LightGBM 5-Fold CV ── print(f"\n{'='*60}") print("LIGHTGBM 5-FOLD CV") print(f"{'='*60}") import lightgbm as lgb from sklearn.model_selection import GroupKFold gkf = GroupKFold(n_splits=5) oof_pred = np.zeros(len(X_train)) fold_scores = [] params = dict( objective='binary', metric='binary_logloss', scale_pos_weight=scale_pos_weight, num_leaves=63, learning_rate=0.05, n_estimators=600, min_child_samples=50, subsample=0.8, colsample_bytree=0.8, n_jobs=-1, verbosity=-1 ) for fold, (tr_idx, va_idx) in enumerate(gkf.split(X_train, y_train, s1_arr)): tr_set = lgb.Dataset(X_train[tr_idx], label=y_train[tr_idx]) va_set = lgb.Dataset(X_train[va_idx], label=y_train[va_idx]) booster = lgb.train(params, tr_set, valid_sets=[va_set], callbacks=[lgb.early_stopping(30, verbose=False)]) va_prob = booster.predict(X_train[va_idx]) oof_pred[va_idx] = va_prob va_s1 = set(s1_arr[va_idx]) f05 = macro_f05(va_s1, s1_arr[va_idx], tgt_arr[va_idx], va_prob, 0.50, gt_lookup) fold_scores.append(f05) print(f" Fold {fold+1}: macro_F0.5 = {f05:.4f}") mean_cv = np.mean(fold_scores) print(f"\n CV Mean Macro F0.5: {mean_cv:.4f} +/- {np.std(fold_scores):.4f}") # ── PHASE 4: Per-country threshold sweep ── print(f"\n{'='*60}") print("PER-COUNTRY THRESHOLD OPTIMIZATION") print(f"{'='*60}") s1_country = train_s1_sample.set_index('entity_id')['country'].to_dict() cand_countries = np.array([s1_country.get(s, 'Unknown') for s in s1_arr]) per_country_threshold = {} for ctry in np.unique(cand_countries): mask = cand_countries == ctry ctry_s1 = set(s1_arr[mask]) best_t, best_sc = 0.50, -1.0 for t in np.arange(0.30, 0.96, 0.02): sc = macro_f05(ctry_s1, s1_arr[mask], tgt_arr[mask], oof_pred[mask], t, gt_lookup) if sc > best_sc: best_sc, best_t = sc, t per_country_threshold[ctry] = round(float(best_t), 2) print(f" {ctry}: threshold={best_t:.2f} F0.5={best_sc:.4f}") default_t = 0.50 print(f" Default threshold: {default_t}") # ── PHASE 5: Train final model ── print(f"\n{'='*60}") print("TRAINING FINAL MODEL ON ALL DATA") print(f"{'='*60}") full_set = lgb.Dataset(X_train, label=y_train) final_model = lgb.train(params, full_set) final_model.save_model(str(CACHE_DIR / 'model.lgb')) print(" Final model saved.") with open(CACHE_DIR / 'meta.json', 'w') as f: json.dump({ 'feature_names': FEATURE_NAMES, 'per_country_threshold': per_country_threshold, 'default_threshold': default_t, 'cv_mean_f05': float(mean_cv), 'blocking_recall': float(recall_ceil), }, f, indent=2) del X_train, y_train, oof_pred, train_cands gc.collect() # ── PHASE 6: TEST INFERENCE ── print(f"\n{'='*60}") print("TEST INFERENCE") print(f"{'='*60}") print("\nLoading test data...") test_s1 = load_tsv("test_source1.tsv") test_s2 = load_tsv("test_source2.tsv") test_s3 = load_tsv("test_source3.tsv") print("\nNormalizing test S1...") test_s1 = normalize_df(test_s1, 'test_s1') all_test_s1_ids = test_s1['entity_id'].tolist() # Run blocking on test set test_cands = run_full_blocking(test_s1, test_s2, test_s3, label="test") # Extract features per country to manage memory print(f"\n{'='*60}") print("FEATURE EXTRACTION (TEST)") print(f"{'='*60}") needed_test_tids = set(test_cands['target_id']) test_s2_norm = normalize_df(test_s2[test_s2['entity_id'].isin(needed_test_tids)], 'test_s2_needed') test_s3_norm = normalize_df(test_s3[test_s3['entity_id'].isin(needed_test_tids)], 'test_s3_needed') X_test = extract_features(test_cands, test_s1, test_s2_norm, test_s3_norm) del test_s2_norm, test_s3_norm gc.collect() # Predict print("\n Running LightGBM prediction...") test_probs = final_model.predict(X_test) # Apply per-country thresholds test_s1_ctry = test_s1.set_index('entity_id')['country'].to_dict() test_s1_ids = test_cands['s1_id'].values test_tgt_ids = test_cands['target_id'].values thresholds = np.array([per_country_threshold.get(test_s1_ctry.get(s, ''), default_t) for s in test_s1_ids]) above = test_probs >= thresholds print(f" Pairs above threshold: {above.sum():,} / {len(above):,} ({above.mean():.1%})") # ── PHASE 7: 1:Many Target Exclusivity ── print(f"\n{'='*60}") print("TARGET EXCLUSIVITY RESOLUTION") print(f"{'='*60}") # Each target -> at most 1 S1 (highest prob); each S1 -> many targets target_best = {} # target_id -> (s1_id, prob) for i in range(len(test_s1_ids)): if not above[i]: continue tid = test_tgt_ids[i] sid = test_s1_ids[i] prob = float(test_probs[i]) if tid not in target_best or prob > target_best[tid][1]: target_best[tid] = (sid, prob) # Build S1 -> list of targets s1_matches = defaultdict(list) for tid, (sid, _) in target_best.items(): s1_matches[sid].append(tid) n_matched_s1 = len(s1_matches) n_singletons = len(all_test_s1_ids) - n_matched_s1 total_links = sum(len(v) for v in s1_matches.values()) print(f" Matched S1 entities: {n_matched_s1:,}") print(f" Singletons: {n_singletons:,}") print(f" Total match links: {total_links:,}") # ── PHASE 8: OUTPUT FILES ── print(f"\n{'='*60}") print("WRITING OUTPUT FILES") print(f"{'='*60}") # matching_results.tsv match_path = OUT_DIR / 'matching_results.tsv' with open(match_path, 'w', encoding='utf-8', newline='') as f: f.write("source1_entity_id\tmatched_entity_ids\n") for sid in all_test_s1_ids: targets = s1_matches.get(sid, []) f.write(f"{sid}\t{','.join(sorted(targets))}\n") print(f" Wrote {match_path}") # candidate_pairs.tsv cand_path = OUT_DIR / 'candidate_pairs.tsv' # Build per-S1 candidate lists from test_cands s1_cand_map = defaultdict(list) for sid, tid in zip(test_cands['s1_id'], test_cands['target_id']): s1_cand_map[sid].append(tid) with open(cand_path, 'w', encoding='utf-8', newline='') as f: f.write("source1_entity_id\tcandidate_entity_ids\n") for sid in all_test_s1_ids: cids = s1_cand_map.get(sid, []) f.write(f"{sid}\t{','.join(cids)}\n") print(f" Wrote {cand_path}") total_time = time.time() - t0_global print(f"\n{'='*60}") print(f"PIPELINE COMPLETE IN {total_time/60:.1f} MINUTES") print(f"{'='*60}") if __name__ == '__main__': main()