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 import ir_datasets from constants import BEIR, IR_DATASETS, LOCAL_DATASETS @st.cache_data def load_local_corpus(corpus_file, columns_to_combine=["title", "text"]): if corpus_file is None: return None did2text = {} id_key = "_id" with corpus_file as f: for idx, line in enumerate(f): uses_bytes = not (type(line) == str) if uses_bytes: if idx == 0 and "doc_id" in line.decode("utf-8"): continue inst = json.loads(line.decode("utf-8")) else: if idx == 0 and "doc_id" in line: continue inst = json.loads(line) all_text = " ".join([inst[col] for col in columns_to_combine if col in inst]) if id_key not in inst: id_key = "doc_id" did2text[inst[id_key]] = { "text": all_text, "title": inst["title"] if "title" in inst else "", } return did2text @st.cache_data def load_local_queries(queries_file): if queries_file is None: return None qid2text = {} id_key = "_id" with queries_file as f: for idx, line in enumerate(f): uses_bytes = not (type(line) == str) if uses_bytes: if idx == 0 and "query_id" in line.decode("utf-8"): continue inst = json.loads(line.decode("utf-8")) else: if idx == 0 and "query_id" in line: continue inst = json.loads(line) if id_key not in inst: id_key = "query_id" qid2text[inst[id_key]] = inst["text"] return qid2text @st.cache_data def load_local_qrels(qrels_file): if qrels_file is None: return None qid2did2label = defaultdict(dict) with qrels_file as f: for idx, line in enumerate(f): uses_bytes = not (type(line) == str) if uses_bytes: if idx == 0 and "qid" in line.decode("utf-8") or "query-id" in line.decode("utf-8"): continue cur_line = line.decode("utf-8") else: if idx == 0 and "qid" in line or "query-id" in line: continue cur_line = line try: qid, _, doc_id, label = cur_line.split() except: qid, doc_id, label = cur_line.split() qid2did2label[str(qid)][str(doc_id)] = int(label) return qid2did2label @st.cache_data 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) all_groups = [] for qid, sub_df in run_pandas.groupby("qid"): sub_df.sort_values(["score", "doc_id"], ascending=[False, False]) sub_df["rank"] = list(range(1, len(sub_df) + 1)) all_groups.append(sub_df) run_pandas = pd.concat(all_groups) return new_run, run_pandas @st.cache_data 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 @st.cache_data(persist="disk") 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") @st.cache_data(persist="disk") def get_ir_datasets(dataset_name: str, input_fields_doc: str = None, input_fields_query: str = None): dataset = ir_datasets.load(dataset_name) queries = {} for qid, query in dataset.queries_iter(): if input_fields_query is None: if type(query) == str: queries[qid] = query else: # get all fields that exist in query all_fields = {field: getattr(query, field) for field in query._fields} # put all fields into a single string queries[qid] = " ".join([str(v) for v in all_fields.values()]) else: all_fields = {field: getattr(query, field) for field in input_fields_query} queries[qid] = " ".join([str(v) for v in all_fields.values()]) corpus = {} for doc in dataset.docs_iter(): if input_fields_doc is None: if type(doc) == str: corpus[doc.doc_id] = {"text": doc} else: # get all fields that exist in query all_fields = {field: getattr(doc, field) for field in doc._fields} corpus[doc.doc_id] = {"text": " ".join([str(v) for v in all_fields.values()])} else: all_fields = {field: getattr(doc, field) for field in input_fields_doc} corpus[doc.doc_id] = {"text": " ".join([str(v) for v in all_fields.values()])} # return corpus, queries, qrels return corpus, queries, dataset.qrels_dict() @st.cache_data(persist="disk") def get_dataset(dataset_name: str, input_fields_doc, input_fields_query): if type(input_fields_doc) == str: input_fields_doc = input_fields_doc.strip().split(",") if type(input_fields_query) == str: input_fields_query = input_fields_query.strip().split(",") if dataset_name == "": return {}, {}, {} if dataset_name in BEIR: return get_beir(dataset_name) elif dataset_name in IR_DATASETS: return get_ir_datasets(dataset_name, input_fields_doc, input_fields_query) elif dataset_name in LOCAL_DATASETS: base_path = f"local_datasets/{dataset_name}" corpus_file = open(f"{base_path}/corpus.jsonl", "r") queries_file = open(f"{base_path}/queries.jsonl", "r") qrels_file = open(f"{base_path}/qrels/test.tsv", "r") return load_local_corpus(corpus_file), load_local_queries(queries_file), load_local_qrels(qrels_file) else: raise NotImplementedError("Dataset not implemented")