chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,399 @@
|
||||
import os
|
||||
import json
|
||||
import logging
|
||||
import datasets
|
||||
from tqdm import tqdm
|
||||
from typing import List, Optional
|
||||
from collections import defaultdict
|
||||
|
||||
from FlagEmbedding.abc.evaluation import AbsEvalDataLoader
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BrightShortEvalDataLoader(AbsEvalDataLoader):
|
||||
"""
|
||||
Data loader class for Bright(short).
|
||||
"""
|
||||
def available_dataset_names(self) -> List[str]:
|
||||
"""
|
||||
Get the available dataset names.
|
||||
|
||||
Returns:
|
||||
List[str]: All the available dataset names.
|
||||
"""
|
||||
return [
|
||||
# StackExchange
|
||||
"biology", "earth_science", "economics", "psychology", "robotics", "stackoverflow", "sustainable_living",
|
||||
# Coding
|
||||
"leetcode", "pony",
|
||||
# Theorem-based
|
||||
"aops", "theoremqa_questions", "theoremqa_theorems"
|
||||
]
|
||||
|
||||
def available_splits(self, dataset_name: str) -> List[str]:
|
||||
"""
|
||||
Get the avaialble splits.
|
||||
|
||||
Args:
|
||||
dataset_name (str): Dataset name.
|
||||
|
||||
Returns:
|
||||
List[str]: All the available splits for the dataset.
|
||||
"""
|
||||
return [
|
||||
# normal splits
|
||||
"examples",
|
||||
# w/ reasoning splits
|
||||
"Gemini-1.0_reason", "claude-3-opus_reason", "gpt4_reason", "grit_reason", "llama3-70b_reason",
|
||||
]
|
||||
|
||||
def _load_remote_corpus(
|
||||
self,
|
||||
dataset_name: str,
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Load the corpus dataset from HF.
|
||||
|
||||
Args:
|
||||
dataset_name (str): Name of the dataset.
|
||||
save_dir (Optional[str], optional): Directory to save the dataset. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: Loaded datasets instance of corpus.
|
||||
"""
|
||||
corpus = datasets.load_dataset(
|
||||
"xlangai/bright", "documents",
|
||||
cache_dir=self.cache_dir,
|
||||
download_mode=self.hf_download_mode
|
||||
)[dataset_name]
|
||||
|
||||
if save_dir is not None:
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
save_path = os.path.join(save_dir, "corpus.jsonl")
|
||||
corpus_dict = {}
|
||||
with open(save_path, "w", encoding="utf-8") as f:
|
||||
for data in tqdm(corpus, desc="Loading and Saving corpus"):
|
||||
docid, text = str(data["id"]), data["content"]
|
||||
_data = {
|
||||
"id": docid,
|
||||
"text": text
|
||||
}
|
||||
corpus_dict[docid] = {"text": text}
|
||||
f.write(json.dumps(_data, ensure_ascii=False) + "\n")
|
||||
logging.info(f"{self.eval_name} {dataset_name} corpus saved to {save_path}")
|
||||
else:
|
||||
corpus_dict = {str(data["id"]): {"text": data["content"]} for data in tqdm(corpus, desc="Loading corpus")}
|
||||
return datasets.DatasetDict(corpus_dict)
|
||||
|
||||
def _load_remote_qrels(
|
||||
self,
|
||||
dataset_name: str,
|
||||
split: str = 'examples',
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Load the qrels from HF.
|
||||
|
||||
Args:
|
||||
dataset_name (str): Name of the dataset.
|
||||
split (str, optional): Split of the dataset. Defaults to ``'examples'``.
|
||||
save_dir (Optional[str], optional): Directory to save the dataset. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: Loaded datasets instance of qrel.
|
||||
"""
|
||||
examples = datasets.load_dataset(
|
||||
"xlangai/bright", split,
|
||||
cache_dir=self.cache_dir,
|
||||
download_mode=self.hf_download_mode
|
||||
)[dataset_name]
|
||||
|
||||
if save_dir is not None:
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
save_path = os.path.join(save_dir, f"{split}_qrels.jsonl")
|
||||
qrels_dict = defaultdict(dict)
|
||||
with open(save_path, "w", encoding="utf-8") as f:
|
||||
for data in tqdm(examples, desc="Loading and Saving qrels"):
|
||||
|
||||
# NOTE: we modify the qid here to distinguish the queries from different splits
|
||||
qid = f'{split}-{data["id"]}'
|
||||
|
||||
for docid in data["gold_ids"]:
|
||||
_data = {
|
||||
"qid": qid,
|
||||
"docid": docid,
|
||||
"relevance": 1
|
||||
}
|
||||
qrels_dict[qid][docid] = 1
|
||||
f.write(json.dumps(_data, ensure_ascii=False) + "\n")
|
||||
|
||||
# NOTE: we record the excluded_ids in qrels with relevance 0 to remove corresponding documents from raw search results. Refer to `searcher.py` for details.
|
||||
for ex_docid in list(set(data["excluded_ids"])):
|
||||
if ex_docid == "N/A":
|
||||
continue
|
||||
assert ex_docid not in qrels_dict[qid], f"{ex_docid} in {qid}"
|
||||
_data = {
|
||||
"qid": qid,
|
||||
"docid": ex_docid,
|
||||
"relevance": 0
|
||||
}
|
||||
qrels_dict[qid][ex_docid] = 0
|
||||
f.write(json.dumps(_data, ensure_ascii=False) + "\n")
|
||||
else:
|
||||
qrels_dict = defaultdict(dict)
|
||||
for data in tqdm(examples, desc="Loading qrels"):
|
||||
|
||||
# NOTE: we modify the qid here to distinguish the queries from different splits
|
||||
qid = f'{split}-{data["id"]}'
|
||||
|
||||
for docid in data["gold_ids"]:
|
||||
qrels_dict[qid][docid] = 1
|
||||
|
||||
# NOTE: we record the excluded_ids in qrels with relevance 0 to remove corresponding documents from raw search results. Refer to `searcher.py` for details.
|
||||
for ex_docid in data["excluded_ids"]:
|
||||
if ex_docid == "N/A":
|
||||
continue
|
||||
assert ex_docid not in qrels_dict[qid], f"{ex_docid} in {qid}"
|
||||
_data = {
|
||||
"qid": qid,
|
||||
"docid": ex_docid,
|
||||
"relevance": 0
|
||||
}
|
||||
qrels_dict[qid][ex_docid] = 0
|
||||
return datasets.DatasetDict(qrels_dict)
|
||||
|
||||
def _load_remote_queries(
|
||||
self,
|
||||
dataset_name: str,
|
||||
split: str = 'examples',
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Load the queries from HF.
|
||||
|
||||
Args:
|
||||
dataset_name (str): Name of the dataset.
|
||||
split (str, optional): Split of the dataset. Defaults to ``'examples'``.
|
||||
save_dir (Optional[str], optional): Directory to save the dataset. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: Loaded datasets instance of queries.
|
||||
"""
|
||||
examples = datasets.load_dataset(
|
||||
"xlangai/bright", split,
|
||||
cache_dir=self.cache_dir,
|
||||
download_mode=self.hf_download_mode
|
||||
)[dataset_name]
|
||||
|
||||
if save_dir is not None:
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
save_path = os.path.join(save_dir, f"{split}_queries.jsonl")
|
||||
queries_dict = {}
|
||||
with open(save_path, "w", encoding="utf-8") as f:
|
||||
for data in tqdm(examples, desc="Loading and Saving queries"):
|
||||
|
||||
# NOTE: we modify the qid here to distinguish the queries from different splits
|
||||
qid, query = f'{split}-{data["id"]}', data["query"]
|
||||
|
||||
_data = {
|
||||
"id": qid,
|
||||
"text": query
|
||||
}
|
||||
queries_dict[qid] = query
|
||||
f.write(json.dumps(_data, ensure_ascii=False) + "\n")
|
||||
else:
|
||||
# NOTE: we modify the qid here to distinguish the queries from different splits
|
||||
queries_dict = {f'{split}-{data["id"]}': data["query"] for data in tqdm(examples, desc="Loading queries")}
|
||||
return datasets.DatasetDict(queries_dict)
|
||||
|
||||
|
||||
class BrightLongEvalDataLoader(AbsEvalDataLoader):
|
||||
"""
|
||||
Data loader class for Bright(long).
|
||||
"""
|
||||
def available_dataset_names(self) -> List[str]:
|
||||
"""
|
||||
Get the available dataset names.
|
||||
|
||||
Returns:
|
||||
List[str]: All the available dataset names.
|
||||
"""
|
||||
return [
|
||||
# StackExchange
|
||||
"biology", "earth_science", "economics", "psychology", "robotics", "stackoverflow", "sustainable_living",
|
||||
# Coding
|
||||
"pony",
|
||||
]
|
||||
|
||||
def available_splits(self, dataset_name: str) -> List[str]:
|
||||
"""
|
||||
Get the avaialble splits.
|
||||
|
||||
Args:
|
||||
dataset_name (str): Dataset name.
|
||||
|
||||
Returns:
|
||||
List[str]: All the available splits for the dataset.
|
||||
"""
|
||||
return [
|
||||
# normal splits
|
||||
"examples",
|
||||
# w/ reasoning splits
|
||||
"Gemini-1.0_reason", "claude-3-opus_reason", "gpt4_reason", "grit_reason", "llama3-70b_reason",
|
||||
]
|
||||
|
||||
def _load_remote_corpus(
|
||||
self,
|
||||
dataset_name: str,
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Load the corpus dataset from HF.
|
||||
|
||||
Args:
|
||||
dataset_name (str): Name of the dataset.
|
||||
save_dir (Optional[str], optional): Directory to save the dataset. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: Loaded datasets instance of corpus.
|
||||
"""
|
||||
corpus = datasets.load_dataset(
|
||||
"xlangai/bright", "long_documents",
|
||||
cache_dir=self.cache_dir,
|
||||
download_mode=self.hf_download_mode
|
||||
)[dataset_name]
|
||||
|
||||
if save_dir is not None:
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
save_path = os.path.join(save_dir, "corpus.jsonl")
|
||||
corpus_dict = {}
|
||||
with open(save_path, "w", encoding="utf-8") as f:
|
||||
for data in tqdm(corpus, desc="Loading and Saving corpus"):
|
||||
docid, text = str(data["id"]), data["content"]
|
||||
_data = {
|
||||
"id": docid,
|
||||
"text": text
|
||||
}
|
||||
corpus_dict[docid] = {"text": text}
|
||||
f.write(json.dumps(_data, ensure_ascii=False) + "\n")
|
||||
logging.info(f"{self.eval_name} {dataset_name} corpus saved to {save_path}")
|
||||
else:
|
||||
corpus_dict = {str(data["id"]): {"text": data["content"]} for data in tqdm(corpus, desc="Loading corpus")}
|
||||
return datasets.DatasetDict(corpus_dict)
|
||||
|
||||
def _load_remote_qrels(
|
||||
self,
|
||||
dataset_name: str,
|
||||
split: str = 'examples',
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Load the qrels from HF.
|
||||
|
||||
Args:
|
||||
dataset_name (str): Name of the dataset.
|
||||
split (str, optional): Split of the dataset. Defaults to ``'examples'``.
|
||||
save_dir (Optional[str], optional): Directory to save the dataset. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: Loaded datasets instance of qrel.
|
||||
"""
|
||||
examples = datasets.load_dataset(
|
||||
"xlangai/bright", split,
|
||||
cache_dir=self.cache_dir,
|
||||
download_mode=self.hf_download_mode
|
||||
)[dataset_name]
|
||||
|
||||
if save_dir is not None:
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
save_path = os.path.join(save_dir, f"{split}_qrels.jsonl")
|
||||
qrels_dict = defaultdict(dict)
|
||||
with open(save_path, "w", encoding="utf-8") as f:
|
||||
for data in tqdm(examples, desc="Loading and Saving qrels"):
|
||||
|
||||
# NOTE: we modify the qid here to distinguish the queries from different splits
|
||||
qid = f'{split}-{data["id"]}'
|
||||
|
||||
for docid in data["gold_ids_long"]:
|
||||
_data = {
|
||||
"qid": qid,
|
||||
"docid": docid,
|
||||
"relevance": 1
|
||||
}
|
||||
qrels_dict[qid][docid] = 1
|
||||
f.write(json.dumps(_data, ensure_ascii=False) + "\n")
|
||||
|
||||
# NOTE: we record the excluded_ids in qrels with relevance 0 to remove corresponding documents from raw search results. Refer to `searcher.py` for details.
|
||||
for ex_docid in list(set(data["excluded_ids"])):
|
||||
if ex_docid == "N/A":
|
||||
continue
|
||||
assert ex_docid not in qrels_dict[qid], f"{ex_docid} in {qid}"
|
||||
_data = {
|
||||
"qid": qid,
|
||||
"docid": ex_docid,
|
||||
"relevance": 0
|
||||
}
|
||||
qrels_dict[qid][ex_docid] = 0
|
||||
f.write(json.dumps(_data, ensure_ascii=False) + "\n")
|
||||
else:
|
||||
qrels_dict = defaultdict(dict)
|
||||
for data in tqdm(examples, desc="Loading qrels"):
|
||||
|
||||
# NOTE: we modify the qid here to distinguish the queries from different splits
|
||||
qid = f'{split}-{data["id"]}'
|
||||
|
||||
for docid in data["gold_ids_long"]:
|
||||
qrels_dict[qid][docid] = 1
|
||||
|
||||
# NOTE: we record the excluded_ids in qrels with relevance 0 to remove corresponding documents from raw search results. Refer to `searcher.py` for details.
|
||||
for ex_docid in data["excluded_ids"]:
|
||||
if ex_docid == "N/A":
|
||||
continue
|
||||
assert ex_docid not in qrels_dict[qid], f"{ex_docid} in {qid}"
|
||||
_data = {
|
||||
"qid": qid,
|
||||
"docid": ex_docid,
|
||||
"relevance": 0
|
||||
}
|
||||
qrels_dict[qid][ex_docid] = 0
|
||||
return datasets.DatasetDict(qrels_dict)
|
||||
|
||||
def _load_remote_queries(
|
||||
self,
|
||||
dataset_name: str,
|
||||
split: str = 'examples',
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Load the queries from HF.
|
||||
|
||||
Args:
|
||||
dataset_name (str): Name of the dataset.
|
||||
split (str, optional): Split of the dataset. Defaults to ``'examples'``.
|
||||
save_dir (Optional[str], optional): Directory to save the dataset. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: Loaded datasets instance of queries.
|
||||
"""
|
||||
examples = datasets.load_dataset(
|
||||
"xlangai/bright", split,
|
||||
cache_dir=self.cache_dir,
|
||||
download_mode=self.hf_download_mode
|
||||
)[dataset_name]
|
||||
|
||||
if save_dir is not None:
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
save_path = os.path.join(save_dir, f"{split}_queries.jsonl")
|
||||
queries_dict = {}
|
||||
with open(save_path, "w", encoding="utf-8") as f:
|
||||
for data in tqdm(examples, desc="Loading and Saving queries"):
|
||||
|
||||
# NOTE: we modify the qid here to distinguish the queries from different splits
|
||||
qid, query = f'{split}-{data["id"]}', data["query"]
|
||||
|
||||
_data = {
|
||||
"id": qid,
|
||||
"text": query
|
||||
}
|
||||
queries_dict[qid] = query
|
||||
f.write(json.dumps(_data, ensure_ascii=False) + "\n")
|
||||
else:
|
||||
# NOTE: we modify the qid here to distinguish the queries from different splits
|
||||
queries_dict = {f'{split}-{data["id"]}': data["query"] for data in tqdm(examples, desc="Loading queries")}
|
||||
return datasets.DatasetDict(queries_dict)
|
||||
Reference in New Issue
Block a user