import streamlit as st import os import pathlib import beir from beir import util from beir.datasets.data_loader import GenericDataLoader import pytrec_eval import pandas as pd from collections import defaultdict import json import copy def load_jsonl(f): did2text = defaultdict(list) sub_did2text = {} for idx, line in enumerate(f): inst = json.loads(line) if "question" in inst: docid = inst["metadata"][0]["passage_id"] if "doc_id" not in inst else inst["doc_id"] did2text[docid].append(inst["question"]) elif "text" in inst: docid = inst["doc_id"] if "doc_id" in inst else inst["did"] did2text[docid].append(inst["text"]) sub_did2text[inst["did"]] = inst["text"] elif "query" in inst: docid = inst["doc_id"] if "doc_id" in inst else inst["did"] did2text[docid].append(inst["query"]) else: breakpoint() raise NotImplementedError("Need to handle this case") return did2text, sub_did2text def get_beir(dataset: str): url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{}.zip".format(dataset) out_dir = os.path.join(pathlib.Path(__file__).parent.absolute(), "datasets") data_path = util.download_and_unzip(url, out_dir) return GenericDataLoader(data_folder=data_path).load(split="test") def load_run(f_run): run = pytrec_eval.parse_run(copy.deepcopy(f_run)) # convert bytes to strings for keys new_run = defaultdict(dict) for key, sub_dict in run.items(): new_run[key.decode("utf-8")] = {k.decode("utf-8"): v for k, v in sub_dict.items()} run_pandas = pd.read_csv(f_run, header=None, index_col=None, sep="\t") run_pandas.columns = ["qid", "generic", "doc_id", "rank", "score", "model"] run_pandas.doc_id = run_pandas.doc_id.astype(str) run_pandas.qid = run_pandas.qid.astype(str) run_pandas["rank"] = run_pandas["rank"].astype(int) run_pandas.score = run_pandas.score.astype(float) # if run_1_alt is not None: # run_1_alt, run_1_alt_sub = load_jsonl(run_1_alt) return new_run, run_pandas with st.sidebar: dataset_name = st.selectbox("Select a dataset in BEIR", ("scifact", "scidocs","trec-covid", "fever", "fiqa", "nfcorpus", "msmarco", "bioasq", "nq", "hotpotqa", "signal1m", "trec-news", "robust04", "arguana", "quora", "climate-fever", "dbpedia-entity", "webis-touche2020", "cqadupstack")) metric_name = st.selectbox("Select a metric", ("recall_10", "recall_5", "ndcg_cut_10")) # sliderbar of how many Top N to choose top_n = st.slider("Top N", 1, 100, 3) x = st.header('Upload a run file') run1_file = st.file_uploader("Choose a file", key="run1") y = st.header("Upload a second run file") run2_file = st.file_uploader("Choose a file", key="run2") incorrect_only = st.checkbox("Show only incorrect instances", value=False) one_better_than_two = st.checkbox("Show only instances where run 1 is better than run 2", value=False) two_better_than_one = st.checkbox("Show only instances where run 2 is better than run 1", value=False) col1, col2 = st.columns([1, 2], gap="medium") incorrect = 0 is_better_run1_count = 0 is_better_run2_count = 0 checkboxes = None with col1: st.title("Instances") if run1_file is not None: print("Running....") corpus, queries, qrels = get_beir(dataset_name) evaluator = pytrec_eval.RelevanceEvaluator( qrels, pytrec_eval.supported_measures) if run1_file is not None: run1, run1_pandas = load_run(run1_file) results1 = evaluator.evaluate(run1) # dict of instance then metrics then values if run2_file is not None: run2, run2_pandas = load_run(run2_file) results2 = evaluator.evaluate(run2) name_of_columns = ["Overview"] + sorted([str(item) for item in set(run1_pandas.qid.tolist())]) checkboxes = [("Overview", st.checkbox("Overview", key=f"0overview"))] st.divider() for idx, item in enumerate(name_of_columns): is_overview = item == "Overview" if is_overview: continue is_incorrect = False is_better_run1 = False is_better_run2 = False run1_score = results1[str(item)][metric_name] if not is_overview else 1 if run2_file is not None: run2_score = results2[str(item)][metric_name] if not is_overview else 1 if not is_overview and run1_score == 0 or run2_score == 0: incorrect += 1 is_incorrect = True if not is_overview and run1_score > run2_score: is_better_run1_count += 1 is_better_run1 = True elif not is_overview and run2_score > run1_score: is_better_run2_count += 1 is_better_run2 = True if not incorrect_only or is_incorrect: if not one_better_than_two or is_better_run1: if not two_better_than_one or is_better_run2: check = st.checkbox(str(item), key=f"{idx}check") st.divider() checkboxes.append((item, check)) else: if not is_overview and run1_score == 0: incorrect += 1 is_incorrect = True if not incorrect_only or is_incorrect: check = st.checkbox(str(item), key=f"{idx}check") st.divider() checkboxes.append((item, check)) with col2: if checkboxes is not None: st.title(f"Information ({len(checkboxes) - 1 if checkboxes else 0}/{len(name_of_columns) - 1})") else: st.title(f"Information") ### Only one run file if run1_file is not None and run2_file is None: for check_idx, (inst_num, checkbox) in enumerate(checkboxes): if checkbox: if inst_num == "Overview": st.header("Overview") st.markdown("TODO: Add overview") else: st.header(f"Instance Number: {inst_num}") st.subheader(f"Query") query_text = queries[str(inst_num)] st.markdown(query_text) st.divider() ## Documents # relevant relevant_docs = list(qrels[str(inst_num)].keys()) doc_texts = [(doc_id, corpus[doc_id]["title"], corpus[doc_id]["text"]) for doc_id in relevant_docs] st.subheader("Relevant Documents") for (docid, title, text) in doc_texts: st.text_area(f"{docid}: {title}", text) # top ranked pred_doc = run1_pandas[run1_pandas.doc_id.isin(relevant_docs)] rank_pred = pred_doc[pred_doc.qid == str(inst_num)]["rank"].tolist() st.subheader("Ranked of Documents") st.markdown(f"Rank: {rank_pred}") st.divider() if st.checkbox('Show top ranked documents'): st.subheader("Top N Ranked Documents") run1_top_n = run1_pandas[run1_pandas.qid == str(inst_num)][:top_n] run1_top_n_docs = [corpus[str(doc_id)] for doc_id in run1_top_n.doc_id.tolist()] for d_idx, doc in enumerate(run1_top_n_docs): st.text_area(f"{run1_top_n['doc_id'].iloc[d_idx]}: {doc['title']}", doc["text"]) st.divider() st.subheader("Score") st.markdown(f"{results1[str(inst_num)][metric_name]}") break ## Both run files available elif run1_file is not None and run2_file is not None: for check_idx, (inst_num, checkbox) in enumerate(checkboxes): if checkbox: if inst_num == "Overview": st.header("Overview") st.markdown("TODO: Add overview") else: st.header(f"Instance Number: {inst_num}") st.subheader(f"Query") query_text = queries[str(inst_num)] st.markdown(query_text) st.divider() ## Documents # relevant relevant_docs = list(qrels[str(inst_num)].keys()) doc_texts = [(doc_id, corpus[doc_id]["title"], corpus[doc_id]["text"]) for doc_id in relevant_docs] st.subheader("Relevant Documents") for (docid, title, text) in doc_texts: st.text_area(f"{docid}: {title}", text) # top ranked pred_doc1 = run1_pandas[run1_pandas.doc_id.isin(relevant_docs)] rank_pred1 = pred_doc1[pred_doc1.qid == str(inst_num)]["rank"].tolist() pred_doc2 = run2_pandas[run2_pandas.doc_id.isin(relevant_docs)] rank_pred2 = pred_doc2[pred_doc2.qid == str(inst_num)]["rank"].tolist() st.subheader("Ranked of Documents") st.markdown(f"Run 1 Rank: {rank_pred1}") st.markdown(f"Run 2 Rank: {rank_pred2}") st.divider() if st.checkbox('Show top ranked documents for Run 1'): st.subheader("Top N Ranked Documents") run1_top_n = run1_pandas[run1_pandas.qid == str(inst_num)][:top_n] run1_top_n_docs = [corpus[str(doc_id)] for doc_id in run1_top_n.doc_id.tolist()] for d_idx, doc in enumerate(run1_top_n_docs): st.text_area(f"{run1_top_n['doc_id'].iloc[d_idx]}: {doc['title']}", doc["text"]) if st.checkbox('Show top ranked documents for Run 2'): st.subheader("Top N Ranked Documents") run2_top_n = run2_pandas[run2_pandas.qid == str(inst_num)][:top_n] run2_top_n_docs = [corpus[str(doc_id)] for doc_id in run2_top_n.doc_id.tolist()] for d_idx, doc in enumerate(run2_top_n_docs): st.text_area(f"{run2_top_n['doc_id'].iloc[d_idx]}: {doc['title']}", doc["text"]) st.divider() st.subheader("Scores") st.markdown(f"Run 1: {results1[str(inst_num)][metric_name]}") st.markdown(f"Run 2: {results2[str(inst_num)][metric_name]}") break