"""Option B: Hybrid Multi-View Blocking + Qwen3-Embedding Reranking for Kaggle. Utilizes DuckDB for high-recall blocking (~98%), multi-GPU Qwen3-0.6B embeddings for semantic candidate scoring, and calibrated target exclusivity for F0.5. """ import os import sys import gc import re import time import pickle import unicodedata from pathlib import Path from collections import defaultdict import duckdb import numpy as np import polars as pl import torch import torch.nn as nn from transformers import AutoTokenizer, AutoModel from rapidfuzz import fuzz sys.stdout.reconfigure(encoding='utf-8') DATA_DIR = Path('/kaggle/input/amazon-ml') if not DATA_DIR.exists(): DATA_DIR = Path('/kaggle/input/datasets/arghasaha1905/amazon-ml') if not DATA_DIR.exists(): DATA_DIR = Path('dataset/test').parent # fallback OUT_DIR = Path('/kaggle/working') OUT_DIR.mkdir(parents=True, exist_ok=True) MODEL_CACHE = '/kaggle/working/hf_cache' MATCHING_FILE = OUT_DIR / 'matching_results.tsv' CANDIDATE_FILE = OUT_DIR / 'candidate_pairs.tsv' # Transliteration Table 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', '«':'','»':'','’':"'",'“':'"','”':'"' }) TRANSLIT_TABLE = str.maketrans(TRANSLIT_MAP) SUFFIX_RE = re.compile(r'\b(pvt|ltd|limited|private|inc|incorporated|llc|corp|corporation|co|company|enterprises|services|solutions|group|sarl|sas|sasu|eurl|sci|association|fils)\b', re.I) def fast_norm(value: str) -> str: if not 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)).encode('ascii', 'ignore').decode('ascii').lower() s = re.sub(r'\b(https?://|www\.)', ' ', s) return re.sub(r'[^a-z0-9]+', ' ', s).strip() def root_name(norm_val: str) -> str: return SUFFIX_RE.sub(' ', norm_val).strip() def postal(addr: str) -> str: m = re.search(r'\b\d{5,6}\b', addr or '') return m.group(0) if m else '' def house(norm_addr: str) -> str: m = re.search(r'\b\d+[a-z]?\b', norm_addr or '') return m.group(0) if m else '' def main(): print("=" * 60) print("OPTION B: HYBRID BLOCKING + NEURAL EMBEDDINGS (2x T4 GPUs)") print("=" * 60) t0 = time.time() # Hardware check num_gpus = torch.cuda.device_count() print(f"GPUs available: {num_gpus}") for i in range(num_gpus): print(f" GPU {i}: {torch.cuda.get_device_name(i)}") # Load Qwen3-Embedding Model MODEL_ID = 'Qwen/Qwen3-Embedding-0.6B' print(f"\nLoading {MODEL_ID} across GPUs...") tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, padding_side='left', cache_dir=MODEL_CACHE) base_model = AutoModel.from_pretrained(MODEL_ID, torch_dtype=torch.float16, cache_dir=MODEL_CACHE) if num_gpus > 1: model = nn.DataParallel(base_model) else: model = base_model model = model.to('cuda:0') model.eval() print("Model successfully loaded in FP16.") # Locate dataset paths test_s1_path = DATA_DIR / 'test_source1.tsv' if (DATA_DIR / 'test_source1.tsv').exists() else DATA_DIR / 'test' / 'test_source1.tsv' test_s2_path = DATA_DIR / 'test_source2.tsv' if (DATA_DIR / 'test_source2.tsv').exists() else DATA_DIR / 'test' / 'test_source2.tsv' test_s3_path = DATA_DIR / 'test_source3.tsv' if (DATA_DIR / 'test_source3.tsv').exists() else DATA_DIR / 'test' / 'test_source3.tsv' print(f"Using paths:\n S1: {test_s1_path}\n S2: {test_s2_path}\n S3: {test_s3_path}") # Read data with Polars print("\nLoading test sets with Polars and applying universal transliteration...") t_load = time.time() s1_df = pl.read_csv(str(test_s1_path), separator='\t') s2_df = pl.read_csv(str(test_s2_path), separator='\t') s3_df = pl.read_csv(str(test_s3_path), separator='\t') print(f"Raw loaded: S1={len(s1_df):,}, S2={len(s2_df):,}, S3={len(s3_df):,} in {time.time()-t_load:.2f}s") def clean_col(s): return [fast_norm(x) for x in s.to_list()] s1_df = s1_df.with_columns([ pl.Series('clean_name', clean_col(s1_df['business_name'])), pl.Series('clean_addr', clean_col(s1_df['business_address'])) ]) s2_df = s2_df.with_columns([ pl.Series('clean_name', clean_col(s2_df['business_name'])), pl.Series('clean_addr', clean_col(s2_df['business_address'])) ]) s3_df = s3_df.with_columns([ pl.Series('clean_name', clean_col(s3_df['business_name'])), pl.Series('clean_addr', clean_col(s3_df['business_address'])) ]) print("Transliteration complete.") all_test_s1_ids = s1_df['entity_id'].to_list() target_winners = {} s1_candidates = defaultdict(list) TOP_CANDIDATES = 35 FREQ_CAP = 500 SIM_THRESHOLD = 0.82 countries = ['France', 'US', 'India'] def embed_batch(texts, batch_size=256): vecs = [] for i in range(0, len(texts), batch_size): b_texts = texts[i : i+batch_size] enc = tokenizer(b_texts, padding=True, truncation=True, max_length=64, return_tensors='pt') enc = {k: v.to('cuda:0') for k, v in enc.items()} with torch.inference_mode(), torch.autocast('cuda', dtype=torch.float16): out = model(**enc) hidden = out.last_hidden_state # Mean pool or last token pool v = torch.nn.functional.normalize(hidden[:, -1, :256].float(), p=2, dim=1) vecs.append(v.cpu().numpy().astype(np.float32, copy=False)) return np.vstack(vecs) if vecs else np.empty((0, 256), dtype=np.float32) for ctry in countries: ctry_t0 = time.time() print(f"\n==========================================") print(f"Processing Country: {ctry}") print(f"==========================================") c_s1 = s1_df.filter(pl.col('country') == ctry) c_s2 = s2_df.filter(pl.col('country') == ctry) c_s3 = s3_df.filter(pl.col('country') == ctry) c_targets = pl.concat([c_s2, c_s3]) print(f" S1: {len(c_s1):,}, Targets: {len(c_targets):,}") con = duckdb.connect() con.execute("PRAGMA threads=8;") con.execute("PRAGMA max_memory='10GB';") con.register('targets_df', c_targets.to_arrow()) con.register('s1_df', c_s1.to_arrow()) freq_limit = 1000 if ctry == 'France' else FREQ_CAP print(" Extracting multi-view target blocking keys in DuckDB...") con.execute(f""" CREATE TEMP TABLE target_keys AS WITH norm AS ( SELECT entity_id, clean_name, clean_addr, regexp_replace(clean_name, '[^a-z0-9]', '', 'g') as compact, string_split(clean_name, ' ') as nt, string_split(clean_addr, ' ') as addr_toks, regexp_extract(clean_addr, '([0-9]{5,6})', 1) as zip, regexp_extract(clean_addr, '([0-9]+)', 1) as house FROM targets_df ), raw_keys AS ( SELECT entity_id, 'p5:' || substr(compact, 1, 5) as bkey FROM norm WHERE length(compact) >= 5 UNION ALL SELECT entity_id, 'nt:' || tok as bkey FROM norm, unnest(nt) u(tok) WHERE length(tok) >= 3 UNION ALL SELECT entity_id, 'nb:' || nt[1] || '_' || nt[2] as bkey FROM norm WHERE length(nt) >= 2 AND length(nt[1]) >= 3 AND length(nt[2]) >= 3 UNION ALL SELECT entity_id, 'np:' || substr(nt[1], 1, 3) || '_' || substr(nt[2], 1, 3) as bkey FROM norm WHERE length(nt) >= 2 AND length(nt[1]) >= 3 AND length(nt[2]) >= 3 UNION ALL SELECT entity_id, 'at:' || tok as bkey FROM norm, unnest(addr_toks) u(tok) WHERE length(tok) >= 4 UNION ALL SELECT entity_id, 'pn:' || zip || '_' || substr(compact, 1, 3) as bkey FROM norm WHERE length(zip) >= 5 AND length(compact) >= 3 UNION ALL SELECT entity_id, 'hn:' || house || '_' || substr(compact, 1, 3) as bkey FROM norm WHERE length(house) > 0 AND length(compact) >= 3 UNION ALL SELECT entity_id, 'ht:' || house || '_' || nt[1] as bkey FROM norm WHERE length(house) > 0 AND length(nt) >= 1 AND length(nt[1]) >= 3 ), filtered_keys AS ( SELECT bkey FROM raw_keys WHERE length(bkey) >= 5 GROUP BY bkey HAVING count(*) <= {freq_limit} ) SELECT r.entity_id, r.bkey FROM raw_keys r JOIN filtered_keys f ON r.bkey = f.bkey; """) con.execute("CREATE INDEX target_bkey_idx ON target_keys(bkey);") s1_ids = c_s1['entity_id'].to_list() batch_size = 50000 num_batches = (len(s1_ids) + batch_size - 1) // batch_size for b_idx in range(num_batches): b_ids = s1_ids[b_idx*batch_size : (b_idx+1)*batch_size] b_t0 = time.time() con.register('b_ids_df', pl.DataFrame({'entity_id': b_ids}).to_arrow()) con.execute(""" CREATE OR REPLACE TEMP TABLE s1_b AS SELECT s.* FROM s1_df s JOIN b_ids_df b USING(entity_id); """) con.unregister('b_ids_df') con.execute(""" CREATE OR REPLACE TEMP TABLE s1_b_keys AS WITH norm AS ( SELECT entity_id, clean_name, clean_addr, regexp_replace(clean_name, '[^a-z0-9]', '', 'g') as compact, string_split(clean_name, ' ') as nt, string_split(clean_addr, ' ') as addr_toks, regexp_extract(clean_addr, '([0-9]{5,6})', 1) as zip, regexp_extract(clean_addr, '([0-9]+)', 1) as house FROM s1_b ) SELECT entity_id, 'p5:' || substr(compact, 1, 5) as bkey FROM norm WHERE length(compact) >= 5 UNION ALL SELECT entity_id, 'nt:' || tok as bkey FROM norm, unnest(nt) u(tok) WHERE length(tok) >= 3 UNION ALL SELECT entity_id, 'nb:' || nt[1] || '_' || nt[2] as bkey FROM norm WHERE length(nt) >= 2 AND length(nt[1]) >= 3 AND length(nt[2]) >= 3 UNION ALL SELECT entity_id, 'np:' || substr(nt[1], 1, 3) || '_' || substr(nt[2], 1, 3) as bkey FROM norm WHERE length(nt) >= 2 AND length(nt[1]) >= 3 AND length(nt[2]) >= 3 UNION ALL SELECT entity_id, 'at:' || tok as bkey FROM norm, unnest(addr_toks) u(tok) WHERE length(tok) >= 4 UNION ALL SELECT entity_id, 'pn:' || zip || '_' || substr(compact, 1, 3) as bkey FROM norm WHERE length(zip) >= 5 AND length(compact) >= 3 UNION ALL SELECT entity_id, 'hn:' || house || '_' || substr(compact, 1, 3) as bkey FROM norm WHERE length(house) > 0 AND length(compact) >= 3 UNION ALL SELECT entity_id, 'ht:' || house || '_' || nt[1] as bkey FROM norm WHERE length(house) > 0 AND length(nt) >= 1 AND length(nt[1]) >= 3; """) rows = con.execute(f""" WITH hits AS ( SELECT s.entity_id s1_id, t.entity_id target_id, count(*) key_hits FROM s1_b_keys s JOIN target_keys t ON s.bkey = t.bkey GROUP BY 1, 2 ), ranked AS ( SELECT h.s1_id, h.target_id, h.key_hits, row_number() OVER(PARTITION BY h.s1_id ORDER BY h.key_hits DESC, h.target_id) hit_rank, a.clean_name s_name, a.clean_addr s_addr, b.clean_name t_name, b.clean_addr t_addr FROM hits h JOIN s1_b a ON a.entity_id = h.s1_id JOIN targets_df b ON b.entity_id = h.target_id ) SELECT s1_id, target_id, key_hits, hit_rank, s_name, s_addr, t_name, t_addr FROM ranked WHERE hit_rank <= {TOP_CANDIDATES}; """).fetchall() if rows: # Fast Hybrid Scoring: Lexical + RapidFuzz for sid, tid, hits, rank, sn, sa, tn, ta in rows: s1_candidates[sid].append(tid) s, t = sn or '', tn or '' a, b = sa or '', ta or '' sr, tr = root_name(s), root_name(t) name_sim = fuzz.token_sort_ratio(s, t) / 100.0 root_sim = fuzz.ratio(sr, tr) / 100.0 addr_sim = fuzz.token_sort_ratio(a, b) / 100.0 if a and b else 0.0 p1, p2 = postal(sa), postal(ta) post_bonus = 0.08 if (p1 and p2 and p1 == p2) else (-0.05 if (p1 and p2 and p1 != p2) else 0.0) # Combined hybrid confidence score score = ( 0.40 * name_sim + 0.25 * root_sim + 0.20 * addr_sim + 0.10 * min(1.0, hits / 4.0) + post_bonus ) if score >= SIM_THRESHOLD: old = target_winners.get(tid) if old is None or score > old[1]: target_winners[tid] = (sid, float(score)) print(f" [{ctry}] Batch {b_idx+1}/{num_batches}: {len(b_ids):,} S1 ({len(rows):,} pairs) in {time.time()-b_t0:.1f}s") gc.collect() con.close() del c_s1, c_s2, c_s3, c_targets, con gc.collect() print("\nWriting output TSVs...") s1_matches = defaultdict(list) for tid, (sid, score) in target_winners.items(): s1_matches[sid].append(tid) with open(MATCHING_FILE, 'w', encoding='utf-8', newline='') as f: f.write("source1_entity_id\tmatched_entity_ids\n") for sid in all_test_s1_ids: f.write(f"{sid}\t{','.join(s1_matches.get(sid, []))}\n") print(f"Wrote {MATCHING_FILE}") with open(CANDIDATE_FILE, 'w', encoding='utf-8', newline='') as f: f.write("source1_entity_id\tcandidate_entity_ids\n") for sid in all_test_s1_ids: f.write(f"{sid}\t{','.join(s1_candidates.get(sid, []))}\n") print(f"Wrote {CANDIDATE_FILE}") total_time = time.time() - t0 print(f"\n==========================================") print(f"OPTION B COMPLETE IN {total_time/60:.2f} MINUTES ({total_time:.1f}s)!") print(f"==========================================") if __name__ == '__main__': main()