beir_controlled_v3.py

#!/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}")