scripts/score_retrieval.py

#!/usr/bin/env python3
"""Compute query-level and aggregate retrieval metrics for labeled JSON. Offline fixture only."""
import argparse, json, math
from pathlib import Path

def score(data,k):
    rows=[]
    for q in data['queries']:
        truth=set(q['relevant_ids']); candidates=q['retrieved_ids'][:k]
        ranks=[j+1 for j,x in enumerate(candidates) if x in truth]
        if len(candidates)!=len(set(candidates)):
            raise ValueError(f"duplicate retrieved IDs: {q['id']}")
        if not truth:
            rows.append({'id':q['id'],'answerable':False,'recall':None,'hit':None,'mrr':None,'ndcg':None});continue
        recall=len(set(candidates)&truth)/len(truth)
        hit=1 if ranks else 0
        mrr=1/min(ranks) if ranks else 0
        dcg=sum(1/math.log2(j+1) for j in ranks)
        idcg=sum(1/math.log2(j+1) for j in range(1,min(k,len(truth))+1))
        rows.append({'id':q['id'],'answerable':True,'recall':recall,'hit':hit,'mrr':mrr,'ndcg':dcg/idcg if idcg else 0})
    valid=[r for r in rows if r['answerable']]
    return {'k':k,'answerable_queries':len(valid),'no_answer_queries':len(rows)-len(valid),
            'mean_recall':sum(x['recall'] for x in valid)/len(valid) if valid else None,
            'mean_hit':sum(x['hit'] for x in valid)/len(valid) if valid else None,
            'mean_mrr':sum(x['mrr'] for x in valid)/len(valid) if valid else None,
            'mean_ndcg':sum(x['ndcg'] for x in valid)/len(valid) if valid else None,
            'per_query':rows}

def main():
    p=argparse.ArgumentParser();p.add_argument('--input',required=True);p.add_argument('--k',type=int,default=3)
    a=p.parse_args()
    if a.k<1: p.error('--k must be >=1')
    data=json.loads(Path(a.input).read_text())
    print(json.dumps(score(data,a.k),indent=2))
if __name__=='__main__': main()