| |
| """ |
| 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 |
|
|
| |
| |
| |
| SAMPLE_N_PER_COUNTRY = 100_000 |
| TOP_K = 50 |
| FREQ_CAP_DEFAULT = 800 |
| FREQ_CAP_FRANCE = 1500 |
| S1_BATCH_SIZE = 50_000 |
|
|
| |
| 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}") |
|
|
| |
| |
| |
| 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 '' |
|
|
| |
| for raw in ["रियल कंसल्टेंट्स प्राइवेट लिमिटेड", "Thérapy", "பாலாஜி"]: |
| out = fast_norm(raw) |
| assert out != '', f"Transliteration FAILED for {raw!r}" |
| print("Transliteration sanity check: PASSED") |
|
|
|
|
| |
| |
| |
| 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") |
|
|
| |
| |
| |
| 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%}" |
|
|
| |
| |
| |
| 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] |
|
|
| |
| 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 |
|
|
|
|
| |
| |
| |
| 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 |
|
|
|
|
| |
| |
| |
| 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...") |
|
|
| |
| 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 |
|
|
|
|
| |
| |
| |
| 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) |
|
|
|
|
| |
| |
| |
| 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)) |
|
|
|
|
| |
| |
| |
| def main(): |
| global train_s1_sample |
|
|
| |
| 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.") |
|
|
| |
| print(f"\n{'='*60}") |
| print("FEATURE EXTRACTION (TRAIN)") |
| print(f"{'='*60}") |
|
|
| |
| 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() |
|
|
| |
| 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}") |
|
|
| |
| 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}") |
|
|
| |
| 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() |
|
|
| |
| 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() |
|
|
| |
| test_cands = run_full_blocking(test_s1, test_s2, test_s3, label="test") |
|
|
| |
| 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() |
|
|
| |
| print("\n Running LightGBM prediction...") |
| test_probs = final_model.predict(X_test) |
|
|
| |
| 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%})") |
|
|
| |
| print(f"\n{'='*60}") |
| print("TARGET EXCLUSIVITY RESOLUTION") |
| print(f"{'='*60}") |
|
|
| |
| target_best = {} |
| 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) |
|
|
| |
| 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:,}") |
|
|
| |
| print(f"\n{'='*60}") |
| print("WRITING OUTPUT FILES") |
| print(f"{'='*60}") |
|
|
| |
| 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}") |
|
|
| |
| cand_path = OUT_DIR / 'candidate_pairs.tsv' |
| |
| 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() |
|
|
