| 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) |
|
|
|
|
| |
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
|
|
