#!/usr/bin/env python3 """Re-run MiniLM + Hadith ablation from saved pool, then report. FIX v3: contents field + sorted baseline + exception wrapper.""" import json, os, time, math, gc, csv, traceback import numpy as np from collections import defaultdict os.environ["OPENAI_API_KEY"] = "dummy" INDEX_DIR = "/workspace/pyserini_index" QUERIES_FILE = "/workspace/fever/queries.jsonl" QRELS_FILE = "/workspace/fever/qrels/test.tsv" POOL_FILE = "/workspace/beir_pool.json" PROGRESS_FILE = "/workspace/beir_controlled_v3_log.txt" RESULTS_FILE = "/workspace/beir_controlled_v3_results.txt" def log(msg): print(msg, flush=True) with open(PROGRESS_FILE, "a") as f: f.write(f"[{time.strftime('%H:%M:%S')}] {msg}\n") t_start = time.time() # ── Load data ── log("Loading queries and qrels...") queries = {} with open(QUERIES_FILE) as f: for line in f: d = json.loads(line) queries[d['_id']] = d['text'] qrels = {} with open(QRELS_FILE) as f: reader = csv.reader(f, delimiter='\t') next(reader) for row in reader: if not row: continue qrels.setdefault(row[0], {})[row[1]] = int(row[2]) # ── Load pool ── log("Loading saved BM25 pool...") with open(POOL_FILE) as f: pool_data = json.load(f) eval_qids = pool_data["qids"] pool_dict = pool_data["pool"] log(f"{len(eval_qids)} queries with pool data") # ── Pyserini for doc text retrieval ── log("Loading Pyserini searcher...") from pyserini.search import SimpleSearcher searcher = SimpleSearcher(INDEX_DIR) searcher.set_bm25(k1=1.2, b=0.75) def ndcg10(ranked_list, gt): dcg = sum((2**gt.get(did,0)-1)/math.log2(k+2) for k,(did,_) in enumerate(ranked_list[:10])) ig = sorted(gt.values(), reverse=True) idcg = sum((2**ig[k]-1)/math.log2(k+2) for k in range(min(10,len(ig)))) return dcg/idcg if idcg > 0 else 0.0 def recall100(ranked_list, gt): if not gt: return 0.0 retrieved = set(did for did,_ in ranked_list[:100]) relevant = set(gt.keys()) return len(retrieved & relevant) / len(relevant) # ── MiniLM ── log("Loading MiniLM...") import torch from transformers import AutoTokenizer, AutoModel torch.set_num_threads(4) model_name = "sentence-transformers/all-MiniLM-L6-v2" tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModel.from_pretrained(model_name).eval() def encode(texts): inputs = tokenizer(texts, padding=True, truncation=True, max_length=128, return_tensors="pt") with torch.no_grad(): emb = model(**inputs).last_hidden_state.mean(dim=1) return emb encode(["warmup"]) log("MiniLM ready") # Configs ABLATIONS = { "bm25": (0, 0, 0, "bm25"), "dense": (0, 0, 0, "dense"), "dense+T": (0.2, 0, 0, "dense"), "dense+S": (0, 0.1, 0, "dense"), "dense+M": (0, 0, 0.2, "dense"), "dense+TS": (0.2, 0.1, 0, "dense"), "dense+MT": (0.2, 0, 0.2, "dense"), "dense+MS": (0, 0.1, 0.2, "dense"), "dense+MTS": (0.2, 0.1, 0.2, "dense"), } CT = [0.3, 0.5, 0.7, 0.8, 0.85, 0.9] TW = [0.1, 0.15, 0.2, 0.3] KSWEEP = [f"tawatur_ct={ct}_tw={tw}" for ct in CT for tw in TW] results = {name: {"ndcg": 0.0, "recall": 0.0} for name in ABLATIONS} for k in KSWEEP: results[k] = {"ndcg": 0.0, "recall": 0.0} qt = 0 errors = 0 t0 = time.time() log(f"Running {len(eval_qids)} queries through MiniLM + Hadith...") for qi, qid in enumerate(eval_qids): try: pool_100 = pool_dict.get(qid, []) if not pool_100: continue pool = [p[0] for p in pool_100] bm25_s = np.array([p[1] for p in pool_100], dtype=float) # Get doc texts — use 'contents' field pool_texts = [] for docid in pool[:100]: doc = searcher.doc(docid) if doc: raw = doc.raw() try: d = json.loads(raw) pool_texts.append(d.get('contents', '')) except: pool_texts.append(raw[:300]) else: pool_texts.append("") if not any(pool_texts): continue # BM25 norm bm25_n = (bm25_s - bm25_s.min()) / (bm25_s.max() - bm25_s.min() + 1e-10) # Encode try: q_emb = encode([queries[qid]]) pool_embs = encode(pool_texts) except Exception as e: errors += 1 if errors <= 5: log(f" ENCODE ERROR qid={qid}: {e}") continue n_pool = len(pool) # Dense scores q_normed = q_emb / (torch.norm(q_emb, dim=1, keepdim=True) + 1e-10) pool_normed = pool_embs / (torch.norm(pool_embs, dim=1, keepdim=True) + 1e-10) dense_scores = (q_normed @ pool_normed.T).cpu().numpy().flatten() dense_n = (dense_scores - dense_scores.min()) / (dense_scores.max() - dense_scores.min() + 1e-10) # Pool similarity matrix sim = (pool_normed @ pool_normed.T).cpu().numpy() # Hadith signals cs = np.sum(sim > 0.7, axis=1).astype(float) csn = (cs - cs.min()) / (cs.max() - cs.min() + 1e-10) iso = cs <= 1 sim_top2 = sim[0][1] if n_pool > 1 else 0 gt = qrels.get(qid, {}) for a_name, (tw, sh, mu, base_type) in ABLATIONS.items(): base = bm25_n if base_type == "bm25" else dense_n if a_name == "bm25" or a_name == "dense": # FIXED: sort by score before evaluating bl_ranked = sorted([(pool[i], float(bm25_n[i] if base_type == 'bm25' else dense_n[i])) for i in range(n_pool)], key=lambda x: -x[1]) results[a_name]["ndcg"] += ndcg10(bl_ranked, gt) results[a_name]["recall"] += recall100(bl_ranked, gt) continue fs = base.copy() if tw > 0: fs += tw * csn if sh > 0: fs -= sh * iso.astype(float) if mu > 0 and n_pool > 2 and sim_top2 > 0.7: fs[(sim[0] > 0.7) | (sim[1] > 0.7)] += mu ranked = sorted([(pool[i], float(fs[i])) for i in range(n_pool)], key=lambda x: -x[1]) results[a_name]["ndcg"] += ndcg10(ranked, gt) results[a_name]["recall"] += recall100(ranked, gt) # Sweep on dense for ct in CT: cs2 = np.sum(sim > ct, axis=1).astype(float) csn2 = (cs2 - cs2.min()) / (cs2.max() - cs2.min() + 1e-10) for tw2 in TW: fs2 = dense_n + tw2 * csn2 ranked = sorted([(pool[i], float(fs2[i])) for i in range(n_pool)], key=lambda x: -x[1]) k = f"tawatur_ct={ct}_tw={tw2}" results[k]["ndcg"] += ndcg10(ranked, gt) results[k]["recall"] += recall100(ranked, gt) qt += 1 del q_emb, pool_embs, pool_normed, sim gc.collect() if qi > 0 and qi % 200 == 0: rate = (qi + 1) / (time.time() - t0) remaining = (len(eval_qids) - qi - 1) / rate if rate > 0 else 0 log(f" {qi}/{len(eval_qids)} @ {rate:.1f}q/s, ~{remaining:.0f}s left") except Exception as e: errors += 1 log(f" ERROR qid={qid}: {e}") log(traceback.format_exc()) if errors > 20: log("TOO MANY ERRORS — aborting") break continue # Finalize for k in results: results[k]["ndcg"] /= max(qt, 1) results[k]["recall"] /= max(qt, 1) dense_bl = results["dense"]["ndcg"] bm25_ndcg_final = results["bm25"]["ndcg"] # Find best sweep best_d = -99 best_c = "" for ct in CT: for tw2 in TW: k = f"tawatur_ct={ct}_tw={tw2}" d = results[k]["ndcg"] - dense_bl if d > best_d: best_d = d best_c = k # ── Report ── elapsed = time.time() - t_start lines = [] lines.append("=" * 70) lines.append(f"BEIR FEVER — Controlled Ablation ({qt} queries, {errors} errors, {elapsed:.0f}s)") lines.append("=" * 70) lines.append(f"{'System':<25} {'NDCG@10':>10} {'Recall@100':>12} {'vs BM25':>12} {'vs Dense':>12}") lines.append("-" * 71) lines.append(f"{'BM25 (k1=1.2,b=0.75)':<25} {bm25_ndcg_final:>10.4f} {results['bm25']['recall']:>12.3f} {'(baseline)':>12}") lines.append(f"{'MiniLM Dense':<25} {dense_bl:>10.4f} {results['dense']['recall']:>12.3f} {dense_bl-bm25_ndcg_final:>+12.4f}") for name in [k for k in ABLATIONS if k not in ("bm25", "dense")]: v = results[name] lines.append(f"{name:<25} {v['ndcg']:>10.4f} {v['recall']:>12.3f} {v['ndcg']-bm25_ndcg_final:>+12.4f} {v['ndcg']-dense_bl:>+12.4f}") if best_c: lines.append(f"{'Best sweep':<25} {results[best_c]['ndcg']:>10.4f} {results[best_c]['recall']:>12.3f} {results[best_c]['ndcg']-bm25_ndcg_final:>+12.4f} {results[best_c]['ndcg']-dense_bl:>+12.4f} [{best_c}]") lines.append("=" * 70) for l in lines: log(l) with open(RESULTS_FILE, "w") as f: f.write("\n".join(lines) + "\n") log(f"\nDone. Results written to {RESULTS_FILE}")