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