import os from typing import Dict, List, Optional, Protocol import pandas as pd import tqdm import ujson from nlp4web_codebase.ir.data_loaders import IRDataset def round_dict(obj: Dict[str, float], ndigits: int = 4) -> Dict[str, float]: return {k: round(v, ndigits=ndigits) for k, v in obj.items()} def sort_dict(obj: Dict[str, float], reverse: bool = True) -> Dict[str, float]: return dict(sorted(obj.items(), key=lambda pair: pair[1], reverse=reverse)) def save_ranking_results( output_dir: str, query_ids: List[str], rankings: List[Dict[str, float]], query_performances_lists: List[Dict[str, float]], cid2tweights_lists: Optional[List[Dict[str, Dict[str, float]]]] = None, ): os.makedirs(output_dir, exist_ok=True) output_path = os.path.join(output_dir, "ranking_results.jsonl") rows = [] for i, (query_id, ranking, query_performances) in enumerate( zip(query_ids, rankings, query_performances_lists) ): row = { "query_id": query_id, "ranking": round_dict(ranking), "query_performances": round_dict(query_performances), "cid2tweights": {}, } if cid2tweights_lists is not None: row["cid2tweights"] = { cid: round_dict(tws) for cid, tws in cid2tweights_lists[i].items() } rows.append(row) pd.DataFrame(rows).to_json( output_path, orient="records", lines=True, ) class TermWeightingFunction(Protocol): def __call__(self, query: str, cid: str) -> Dict[str, float]: ... def compare( dataset: IRDataset, results_path1: str, results_path2: str, output_dir: str, main_metric: str = "recip_rank", system1: Optional[str] = None, system2: Optional[str] = None, term_weighting_fn1: Optional[TermWeightingFunction] = None, term_weighting_fn2: Optional[TermWeightingFunction] = None, ) -> None: os.makedirs(output_dir, exist_ok=True) df1 = pd.read_json(results_path1, orient="records", lines=True) df2 = pd.read_json(results_path2, orient="records", lines=True) assert len(df1) == len(df2) all_qrels = {} for split in dataset.split2qrels: all_qrels.update(dataset.get_qrels_dict(split)) qid2query = {query.query_id: query for query in dataset.queries} cid2doc = {doc.collection_id: doc for doc in dataset.corpus} diff_col = f"{main_metric}:qp1-qp2" merged = pd.merge(df1, df2, on="query_id", how="outer") rows = [] for _, example in tqdm.tqdm(merged.iterrows(), desc="Comparing", total=len(merged)): docs = {cid: cid2doc[cid].text for cid in dict(example["ranking_x"])} docs.update({cid: cid2doc[cid].text for cid in dict(example["ranking_y"])}) query_id = example["query_id"] row = { "query_id": query_id, "query": qid2query[query_id].text, diff_col: example["query_performances_x"][main_metric] - example["query_performances_y"][main_metric], "ranking1": ujson.dumps(example["ranking_x"], indent=4), "ranking2": ujson.dumps(example["ranking_y"], indent=4), "docs": ujson.dumps(docs, indent=4), "query_performances1": ujson.dumps( example["query_performances_x"], indent=4 ), "query_performances2": ujson.dumps( example["query_performances_y"], indent=4 ), "qrels": ujson.dumps(all_qrels[query_id], indent=4), } if term_weighting_fn1 is not None and term_weighting_fn2 is not None: all_cids = set(example["ranking_x"]) | set(example["ranking_y"]) cid2tweights1 = {} cid2tweights2 = {} ranking1 = {} ranking2 = {} for cid in all_cids: tweights1 = term_weighting_fn1(query=qid2query[query_id].text, cid=cid) tweights2 = term_weighting_fn2(query=qid2query[query_id].text, cid=cid) ranking1[cid] = sum(tweights1.values()) ranking2[cid] = sum(tweights2.values()) cid2tweights1[cid] = tweights1 cid2tweights2[cid] = tweights2 ranking1 = sort_dict(ranking1) ranking2 = sort_dict(ranking2) row["ranking1"] = ujson.dumps(ranking1, indent=4) row["ranking2"] = ujson.dumps(ranking2, indent=4) cid2tweights1 = {cid: cid2tweights1[cid] for cid in ranking1} cid2tweights2 = {cid: cid2tweights2[cid] for cid in ranking2} row["cid2tweights1"] = ujson.dumps(cid2tweights1, indent=4) row["cid2tweights2"] = ujson.dumps(cid2tweights2, indent=4) rows.append(row) table = pd.DataFrame(rows).sort_values(by=diff_col, ascending=False) output_path = os.path.join(output_dir, f"compare-{system1}_vs_{system2}.tsv") table.to_csv(output_path, sep="\t", index=False) # if __name__ == "__main__": # # python -m lecture2.bm25.analysis # from nlp4web_codebase.ir.data_loaders.sciq import load_sciq # from lecture2.bm25.bm25_retriever import BM25Retriever # from lecture2.bm25.tfidf_retriever import TFIDFRetriever # import numpy as np # sciq = load_sciq() # system1 = "bm25" # system2 = "tfidf" # results_path1 = f"output/sciq-{system1}/results/ranking_results.jsonl" # results_path2 = f"output/sciq-{system2}/results/ranking_results.jsonl" # index_dir1 = f"output/sciq-{system1}" # index_dir2 = f"output/sciq-{system2}" # compare( # dataset=sciq, # results_path1=results_path1, # results_path2=results_path2, # output_dir=f"output/sciq-{system1}_vs_{system2}", # system1=system1, # system2=system2, # term_weighting_fn1=BM25Retriever(index_dir1).get_term_weights, # term_weighting_fn2=TFIDFRetriever(index_dir2).get_term_weights, # ) # # bias on #shared_terms of TFIDF: # df1 = pd.read_json(results_path1, orient="records", lines=True) # df2 = pd.read_json(results_path2, orient="records", lines=True) # merged = pd.merge(df1, df2, on="query_id", how="outer") # nterms1 = [] # nterms2 = [] # for _, row in merged.iterrows(): # nterms1.append(len(list(dict(row["cid2tweights_x"]).values())[0])) # nterms2.append(len(list(dict(row["cid2tweights_y"]).values())[0])) # percentiles = (5, 25, 50, 75, 95) # print(system1, np.percentile(nterms1, percentiles), np.mean(nterms1).round(2)) # print(system2, np.percentile(nterms2, percentiles), np.mean(nterms2).round(2)) # # bm25 [ 3. 4. 5. 7. 11.] 5.64 # # tfidf [1. 2. 3. 5. 9.] 3.58