nlp4web_codebase/ir/analysis.py

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