| """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 |
|
|
| 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' |
|
|
| |
| 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() |
|
|
| |
| 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)}") |
|
|
| |
| 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.") |
|
|
| |
| 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}") |
|
|
| |
| 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 |
| |
| 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: |
| |
| 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) |
|
|
| |
| 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() |
|
|
