chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,17 @@
|
||||
from .arguments import AbsEvalArgs, AbsEvalModelArgs
|
||||
from .evaluator import AbsEvaluator
|
||||
from .data_loader import AbsEvalDataLoader
|
||||
from .searcher import EvalRetriever, EvalDenseRetriever, EvalReranker
|
||||
from .runner import AbsEvalRunner
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AbsEvalArgs",
|
||||
"AbsEvalModelArgs",
|
||||
"AbsEvaluator",
|
||||
"AbsEvalDataLoader",
|
||||
"EvalRetriever",
|
||||
"EvalDenseRetriever",
|
||||
"EvalReranker",
|
||||
"AbsEvalRunner",
|
||||
]
|
||||
@@ -0,0 +1,196 @@
|
||||
"""
|
||||
Adapted from https://github.com/AIR-Bench/AIR-Bench/blob/0.1.0/air_benchmark/evaluation_utils/evaluation_arguments.py
|
||||
"""
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class AbsEvalArgs:
|
||||
"""
|
||||
Base class for evaluation arguments.
|
||||
"""
|
||||
eval_name: str = field(
|
||||
default=None,
|
||||
metadata={"help": "The name of the evaluation task, such as msmarco, beir, miracl, etc."}
|
||||
)
|
||||
dataset_dir: Optional[str] = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": "1) If you want to perform evaluation on your own dataset, you can provide the path to the dataset directory (must exists in local). "
|
||||
"The dataset directory should contain the following files: corpus.jsonl, <split>_queries.jsonl, <split>_qrels.jsonl, or contain multiple directories, each of which contains the following files: corpus.jsonl, <split>_queries.jsonl, <split>_qrels.jsonl."
|
||||
"2) If you want to perform evaluation on the datasets we provide evaluation APIs for, you can provide the path to saving the downloaded dataset. If you provide None, the dataset will be only downloaded to the cache directory."
|
||||
}
|
||||
)
|
||||
force_redownload: bool = field(
|
||||
default=False, metadata={"help": "Whether to force redownload the dataset. This is useful when you load dataset from remote and want to update the dataset."}
|
||||
)
|
||||
dataset_names: Optional[str] = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": "The names of the datasets to evaluate. Default: None. If None, all available datasets will be evaluated. The name can be a specific dataset name (BEIR), a specific language (MIRACL), etc.",
|
||||
"nargs": "+"
|
||||
}
|
||||
)
|
||||
splits: str = field(
|
||||
default="test",
|
||||
metadata={"help": "Splits to evaluate. Default: test", "nargs": "+"}
|
||||
)
|
||||
corpus_embd_save_dir: str = field(
|
||||
default=None, metadata={"help": "Path to save corpus embeddings. If None, embeddings are not saved."}
|
||||
)
|
||||
output_dir: str = field(
|
||||
default="./search_results", metadata={"help": "Path to save results."}
|
||||
)
|
||||
search_top_k: int = field(
|
||||
default=1000, metadata={"help": "Top k for retrieving."}
|
||||
)
|
||||
rerank_top_k: int = field(default=100, metadata={"help": "Top k for reranking."})
|
||||
cache_path: str = field(
|
||||
default=None, metadata={"help": "Cache directory for loading datasets."}
|
||||
)
|
||||
token: str = field(
|
||||
default_factory=lambda: os.getenv('HF_TOKEN', None),
|
||||
metadata={"help": "The token to use when accessing the model."}
|
||||
)
|
||||
overwrite: bool = field(
|
||||
default=False, metadata={"help": "whether to overwrite evaluation results"}
|
||||
)
|
||||
ignore_identical_ids: bool = field(
|
||||
default=False, metadata={"help": "whether to ignore identical ids in search results"}
|
||||
)
|
||||
# ================ for evaluation ===============
|
||||
k_values: int = field(
|
||||
default_factory=lambda: [1, 3, 5, 10, 100, 1000],
|
||||
metadata={"help": "k values for evaluation. Default: [1, 3, 5, 10, 100, 1000]", "nargs": "+"}
|
||||
)
|
||||
eval_output_method: str = field(
|
||||
default="markdown",
|
||||
metadata={"help": "The output method for evaluation results. Available methods: ['json', 'markdown']. Default: markdown.", "choices": ["json", "markdown"]}
|
||||
)
|
||||
eval_output_path: str = field(
|
||||
default="./eval_results.md", metadata={"help": "The path to save evaluation results."}
|
||||
)
|
||||
eval_metrics: str = field(
|
||||
default_factory=lambda: ["ndcg_at_10", "recall_at_10"],
|
||||
metadata={"help": "The metrics to evaluate. Default: ['ndcg_at_10', 'recall_at_10']", "nargs": "+"}
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AbsEvalModelArgs:
|
||||
"""
|
||||
Base class for model arguments during evaluation.
|
||||
"""
|
||||
embedder_name_or_path: str = field(
|
||||
metadata={"help": "The embedder name or path.", "required": True}
|
||||
)
|
||||
embedder_model_class: Optional[str] = field(
|
||||
default=None, metadata={"help": "The embedder model class. Available classes: ['encoder-only-base', 'encoder-only-m3', 'decoder-only-base', 'decoder-only-icl', 'decoder-only-pseudo_moe']. Default: None. For the custom model, you need to specifiy the model class.", "choices": ["encoder-only-base", "encoder-only-m3", "decoder-only-base", "decoder-only-icl", "decoder-only-pseudo_moe"]}
|
||||
)
|
||||
normalize_embeddings: bool = field(
|
||||
default=True, metadata={"help": "whether to normalize the embeddings"}
|
||||
)
|
||||
pooling_method: Optional[str] = field(
|
||||
default=None, metadata={"help": "The pooling method fot the embedder."}
|
||||
)
|
||||
use_fp16: bool = field(
|
||||
default=True, metadata={"help": "whether to use fp16 for inference"}
|
||||
)
|
||||
devices: Optional[str] = field(
|
||||
default=None, metadata={"help": "Devices to use for inference.", "nargs": "+"}
|
||||
)
|
||||
query_instruction_for_retrieval: Optional[str] = field(
|
||||
default=None, metadata={"help": "Instruction for query"}
|
||||
)
|
||||
query_instruction_format_for_retrieval: str = field(
|
||||
default="{}{}", metadata={"help": "Format for query instruction"}
|
||||
)
|
||||
examples_for_task: Optional[str] = field(
|
||||
default=None, metadata={"help": "Examples for task"}
|
||||
)
|
||||
examples_instruction_format: str = field(
|
||||
default="{}{}", metadata={"help": "Format for examples instruction"}
|
||||
)
|
||||
trust_remote_code: bool = field(
|
||||
default=False, metadata={"help": "Trust remote code"}
|
||||
)
|
||||
reranker_name_or_path: Optional[str] = field(
|
||||
default=None, metadata={"help": "The reranker name or path."}
|
||||
)
|
||||
reranker_model_class: Optional[str] = field(
|
||||
default=None, metadata={"help": "The reranker model class. Available classes: ['encoder-only-base', 'decoder-only-base', 'decoder-only-layerwise', 'decoder-only-lightweight']. Default: None. For the custom model, you need to specify the model class.", "choices": ["encoder-only-base", "decoder-only-base", "decoder-only-layerwise", "decoder-only-lightweight"]}
|
||||
)
|
||||
reranker_peft_path: Optional[str] = field(
|
||||
default=None, metadata={"help": "The reranker peft path."}
|
||||
)
|
||||
use_bf16: bool = field(
|
||||
default=False, metadata={"help": "whether to use bf16 for inference"}
|
||||
)
|
||||
query_instruction_for_rerank: Optional[str] = field(
|
||||
default=None, metadata={"help": "Instruction for query"}
|
||||
)
|
||||
query_instruction_format_for_rerank: str = field(
|
||||
default="{}{}", metadata={"help": "Format for query instruction"}
|
||||
)
|
||||
passage_instruction_for_rerank: Optional[str] = field(
|
||||
default=None, metadata={"help": "Instruction for passage"}
|
||||
)
|
||||
passage_instruction_format_for_rerank: str = field(
|
||||
default="{}{}", metadata={"help": "Format for passage instruction"}
|
||||
)
|
||||
cache_dir: str = field(
|
||||
default=None, metadata={"help": "Cache directory for models."}
|
||||
)
|
||||
domain_for_pseudo_moe: Optional[str] = field(
|
||||
default=None, metadata={"help": "Domain used by decoder-only-pseudo_moe model, e.g. general/coding/reasoning."}
|
||||
)
|
||||
# ================ for inference ===============
|
||||
embedder_batch_size: int = field(
|
||||
default=3000, metadata={"help": "Batch size for inference."}
|
||||
)
|
||||
reranker_batch_size: int = field(
|
||||
default=3000, metadata={"help": "Batch size for inference."}
|
||||
)
|
||||
embedder_query_max_length: int = field(
|
||||
default=512, metadata={"help": "Max length for query."}
|
||||
)
|
||||
embedder_passage_max_length: int = field(
|
||||
default=512, metadata={"help": "Max length for passage."}
|
||||
)
|
||||
truncate_dim: Optional[int] = field(
|
||||
default=None, metadata={"help": "The dimension to truncate embeddings to. Useful for Matryoshka Representation Learning models. If None, no truncation is performed."}
|
||||
)
|
||||
reranker_query_max_length: Optional[int] = field(
|
||||
default=None, metadata={"help": "Max length for reranking."}
|
||||
)
|
||||
reranker_max_length: int = field(
|
||||
default=512, metadata={"help": "Max length for reranking."}
|
||||
)
|
||||
normalize: bool = field(
|
||||
default=False, metadata={"help": "whether to normalize the reranking scores"}
|
||||
)
|
||||
prompt: Optional[str] = field(
|
||||
default=None, metadata={"help": "The prompt for the reranker."}
|
||||
)
|
||||
cutoff_layers: List[int] = field(
|
||||
default=None, metadata={"help": "The output layers of layerwise/lightweight reranker."}
|
||||
)
|
||||
compress_ratio: int = field(
|
||||
default=1, metadata={"help": "The compress ratio of lightweight reranker."}
|
||||
)
|
||||
compress_layers: Optional[int] = field(
|
||||
default=None, metadata={"help": "The compress layers of lightweight reranker.", "nargs": "+"}
|
||||
)
|
||||
|
||||
def __post_init__(self):
|
||||
# replace "\\n" with "\n"
|
||||
if "\\n" in self.query_instruction_format_for_retrieval:
|
||||
self.query_instruction_format_for_retrieval = self.query_instruction_format_for_retrieval.replace("\\n", "\n")
|
||||
if "\\n" in self.examples_instruction_format:
|
||||
self.examples_instruction_format = self.examples_instruction_format.replace("\\n", "\n")
|
||||
if "\\n" in self.query_instruction_format_for_rerank:
|
||||
self.query_instruction_format_for_rerank = self.query_instruction_format_for_rerank.replace("\\n", "\n")
|
||||
if "\\n" in self.passage_instruction_format_for_rerank:
|
||||
self.passage_instruction_format_for_rerank = self.passage_instruction_format_for_rerank.replace("\\n", "\n")
|
||||
@@ -0,0 +1,423 @@
|
||||
"""
|
||||
Adapted from https://github.com/AIR-Bench/AIR-Bench/blob/0.1.0/air_benchmark/evaluation_utils/data_loader.py
|
||||
"""
|
||||
import os
|
||||
import logging
|
||||
import datasets
|
||||
import subprocess
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Optional, Union
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsEvalDataLoader(ABC):
|
||||
"""
|
||||
Base class of data loader for evaluation.
|
||||
|
||||
Args:
|
||||
eval_name (str): The experiment name of current evaluation.
|
||||
dataset_dir (str, optional): path to the datasets. Defaults to ``None``.
|
||||
cache_dir (str, optional): Path to HuggingFace cache directory. Defaults to ``None``.
|
||||
token (str, optional): HF_TOKEN to access the private datasets/models in HF. Defaults to ``None``.
|
||||
force_redownload: If True, will force redownload the dataset to cover the local dataset. Defaults to ``False``.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
eval_name: str,
|
||||
dataset_dir: Optional[str] = None,
|
||||
cache_dir: Optional[str] = None,
|
||||
token: Optional[str] = None,
|
||||
force_redownload: bool = False
|
||||
):
|
||||
self.eval_name = eval_name
|
||||
self.dataset_dir = dataset_dir
|
||||
if cache_dir is None:
|
||||
cache_dir = os.getenv('HF_HUB_CACHE', '~/.cache/huggingface/hub')
|
||||
self.cache_dir = os.path.join(cache_dir, eval_name)
|
||||
self.token = token
|
||||
self.force_redownload = force_redownload
|
||||
self.hf_download_mode = None if not force_redownload else "force_redownload"
|
||||
|
||||
def available_dataset_names(self) -> List[str]:
|
||||
"""
|
||||
Returns: List[str]: Available dataset names.
|
||||
"""
|
||||
return []
|
||||
|
||||
@abstractmethod
|
||||
def available_splits(self, dataset_name: Optional[str] = None) -> List[str]:
|
||||
"""
|
||||
Returns: List[str]: Available splits in the dataset.
|
||||
"""
|
||||
pass
|
||||
|
||||
def check_dataset_names(self, dataset_names: Union[str, List[str]]) -> List[str]:
|
||||
"""Check the validity of dataset names
|
||||
|
||||
Args:
|
||||
dataset_names (Union[str, List[str]]): a dataset name (str) or a list of dataset names (List[str])
|
||||
|
||||
Raises:
|
||||
ValueError
|
||||
|
||||
Returns:
|
||||
List[str]: List of valid dataset names.
|
||||
"""
|
||||
available_dataset_names = self.available_dataset_names()
|
||||
if isinstance(dataset_names, str):
|
||||
dataset_names = [dataset_names]
|
||||
|
||||
for dataset_name in dataset_names:
|
||||
if dataset_name not in available_dataset_names:
|
||||
raise ValueError(f"Dataset name '{dataset_name}' not found in the dataset. Available dataset names: {available_dataset_names}")
|
||||
return dataset_names
|
||||
|
||||
def check_splits(self, splits: Union[str, List[str]], dataset_name: Optional[str] = None) -> List[str]:
|
||||
"""Check whether the splits are available in the dataset.
|
||||
|
||||
Args:
|
||||
splits (Union[str, List[str]]): Splits to check.
|
||||
dataset_name (Optional[str], optional): Name of dataset to check. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
List[str]: The available splits.
|
||||
"""
|
||||
available_splits = self.available_splits(dataset_name=dataset_name)
|
||||
if isinstance(splits, str):
|
||||
splits = [splits]
|
||||
checked_splits = []
|
||||
for split in splits:
|
||||
if split not in available_splits:
|
||||
logger.warning(f"Split '{split}' not found in the dataset. Removing it from the list.")
|
||||
else:
|
||||
checked_splits.append(split)
|
||||
return checked_splits
|
||||
|
||||
def load_corpus(self, dataset_name: Optional[str] = None) -> datasets.DatasetDict:
|
||||
"""Load the corpus from the dataset.
|
||||
|
||||
Args:
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: A dict of corpus with id as key, title and text as value.
|
||||
"""
|
||||
if self.dataset_dir is not None:
|
||||
if dataset_name is None:
|
||||
save_dir = self.dataset_dir
|
||||
else:
|
||||
save_dir = os.path.join(self.dataset_dir, dataset_name)
|
||||
return self._load_local_corpus(save_dir, dataset_name=dataset_name)
|
||||
else:
|
||||
return self._load_remote_corpus(dataset_name=dataset_name)
|
||||
|
||||
def load_qrels(self, dataset_name: Optional[str] = None, split: str = 'test') -> datasets.DatasetDict:
|
||||
"""Load the qrels from the dataset.
|
||||
|
||||
Args:
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
split (str, optional): The split to load relevance from. Defaults to ``'test'``.
|
||||
|
||||
Raises:
|
||||
ValueError
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: A dict of relevance of query and document.
|
||||
"""
|
||||
if self.dataset_dir is not None:
|
||||
if dataset_name is None:
|
||||
save_dir = self.dataset_dir
|
||||
else:
|
||||
checked_dataset_names = self.check_dataset_names(dataset_name)
|
||||
if len(checked_dataset_names) == 0:
|
||||
raise ValueError(f"Dataset name {dataset_name} not found in the dataset.")
|
||||
dataset_name = checked_dataset_names[0]
|
||||
|
||||
save_dir = os.path.join(self.dataset_dir, dataset_name)
|
||||
|
||||
return self._load_local_qrels(save_dir, dataset_name=dataset_name, split=split)
|
||||
else:
|
||||
return self._load_remote_qrels(dataset_name=dataset_name, split=split)
|
||||
|
||||
def load_queries(self, dataset_name: Optional[str] = None, split: str = 'test') -> datasets.DatasetDict:
|
||||
"""Load the queries from the dataset.
|
||||
|
||||
Args:
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
split (str, optional): The split to load queries from. Defaults to ``'test'``.
|
||||
|
||||
Raises:
|
||||
ValueError
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: A dict of queries with id as key, query text as value.
|
||||
"""
|
||||
if self.dataset_dir is not None:
|
||||
if dataset_name is None:
|
||||
save_dir = self.dataset_dir
|
||||
else:
|
||||
checked_dataset_names = self.check_dataset_names(dataset_name)
|
||||
if len(checked_dataset_names) == 0:
|
||||
raise ValueError(f"Dataset name {dataset_name} not found in the dataset.")
|
||||
dataset_name = checked_dataset_names[0]
|
||||
|
||||
save_dir = os.path.join(self.dataset_dir, dataset_name)
|
||||
|
||||
return self._load_local_queries(save_dir, dataset_name=dataset_name, split=split)
|
||||
else:
|
||||
return self._load_remote_queries(dataset_name=dataset_name, split=split)
|
||||
|
||||
def _load_remote_corpus(
|
||||
self,
|
||||
dataset_name: Optional[str] = None,
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Abstract method to load corpus from remote dataset, to be overrode in child class.
|
||||
|
||||
Args:
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
save_dir (Optional[str], optional): Path to save the new downloaded corpus. Defaults to ``None``.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Loading remote corpus is not implemented.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: A dict of corpus with id as key, title and text as value.
|
||||
"""
|
||||
raise NotImplementedError("Loading remote corpus is not implemented.")
|
||||
|
||||
def _load_remote_qrels(
|
||||
self,
|
||||
dataset_name: Optional[str] = None,
|
||||
split: str = 'test',
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Abstract method to load relevance from remote dataset, to be overrode in child class.
|
||||
|
||||
Args:
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
split (str, optional): Split to load from the remote dataset. Defaults to ``'test'``.
|
||||
save_dir (Optional[str], optional): Path to save the new downloaded relevance. Defaults to ``None``.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Loading remote qrels is not implemented.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: A dict of relevance of query and document.
|
||||
"""
|
||||
raise NotImplementedError("Loading remote qrels is not implemented.")
|
||||
|
||||
def _load_remote_queries(
|
||||
self,
|
||||
dataset_name: Optional[str] = None,
|
||||
split: str = 'test',
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Abstract method to load queries from remote dataset, to be overrode in child class.
|
||||
|
||||
Args:
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
split (str, optional): Split to load from the remote dataset. Defaults to ``'test'``.
|
||||
save_dir (Optional[str], optional): Path to save the new downloaded queries. Defaults to ``None``.
|
||||
|
||||
Raises:
|
||||
NotImplementedError
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: A dict of queries with id as key, query text as value.
|
||||
"""
|
||||
raise NotImplementedError("Loading remote queries is not implemented.")
|
||||
|
||||
def _load_local_corpus(self, save_dir: str, dataset_name: Optional[str] = None) -> datasets.DatasetDict:
|
||||
"""Load corpus from local dataset.
|
||||
|
||||
Args:
|
||||
save_dir (str): Path to save the loaded corpus.
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: A dict of corpus with id as key, title and text as value.
|
||||
"""
|
||||
corpus_path = os.path.join(save_dir, 'corpus.jsonl')
|
||||
if self.force_redownload or not os.path.exists(corpus_path):
|
||||
logger.warning(f"Corpus not found in {corpus_path}. Trying to download the corpus from the remote and save it to {save_dir}.")
|
||||
return self._load_remote_corpus(dataset_name=dataset_name, save_dir=save_dir)
|
||||
else:
|
||||
corpus_data = datasets.load_dataset('json', data_files=corpus_path, cache_dir=self.cache_dir)['train']
|
||||
|
||||
corpus = {}
|
||||
for e in corpus_data:
|
||||
corpus[e['id']] = {'title': e.get('title', ""), 'text': e['text']}
|
||||
|
||||
return datasets.DatasetDict(corpus)
|
||||
|
||||
def _load_local_qrels(self, save_dir: str, dataset_name: Optional[str] = None, split: str = 'test') -> datasets.DatasetDict:
|
||||
"""Load relevance from local dataset.
|
||||
|
||||
Args:
|
||||
save_dir (str): Path to save the loaded relevance.
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
split (str, optional): Split to load from the local dataset. Defaults to ``'test'``.
|
||||
|
||||
Raises:
|
||||
ValueError
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: A dict of relevance of query and document.
|
||||
"""
|
||||
checked_split = self.check_splits(split, dataset_name=dataset_name)
|
||||
if len(checked_split) == 0:
|
||||
raise ValueError(f"Split {split} not found in the dataset.")
|
||||
split = checked_split[0]
|
||||
|
||||
qrels_path = os.path.join(save_dir, f"{split}_qrels.jsonl")
|
||||
if self.force_redownload or not os.path.exists(qrels_path):
|
||||
logger.warning(f"Qrels not found in {qrels_path}. Trying to download the qrels from the remote and save it to {save_dir}.")
|
||||
return self._load_remote_qrels(dataset_name=dataset_name, split=split, save_dir=save_dir)
|
||||
else:
|
||||
qrels_data = datasets.load_dataset('json', data_files=qrels_path, cache_dir=self.cache_dir)['train']
|
||||
|
||||
qrels = {}
|
||||
for data in qrels_data:
|
||||
qid = data['qid']
|
||||
if qid not in qrels:
|
||||
qrels[qid] = {}
|
||||
qrels[qid][data['docid']] = data['relevance']
|
||||
|
||||
return datasets.DatasetDict(qrels)
|
||||
|
||||
def _load_local_queries(self, save_dir: str, dataset_name: Optional[str] = None, split: str = 'test') -> datasets.DatasetDict:
|
||||
"""Load queries from local dataset.
|
||||
|
||||
Args:
|
||||
save_dir (str): Path to save the loaded queries.
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
split (str, optional): Split to load from the local dataset. Defaults to ``'test'``.
|
||||
|
||||
Raises:
|
||||
ValueError
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: A dict of queries with id as key, query text as value.
|
||||
"""
|
||||
checked_split = self.check_splits(split, dataset_name=dataset_name)
|
||||
if len(checked_split) == 0:
|
||||
raise ValueError(f"Split {split} not found in the dataset.")
|
||||
split = checked_split[0]
|
||||
|
||||
queries_path = os.path.join(save_dir, f"{split}_queries.jsonl")
|
||||
if self.force_redownload or not os.path.exists(queries_path):
|
||||
logger.warning(f"Queries not found in {queries_path}. Trying to download the queries from the remote and save it to {save_dir}.")
|
||||
return self._load_remote_queries(dataset_name=dataset_name, split=split, save_dir=save_dir)
|
||||
else:
|
||||
queries_data = datasets.load_dataset('json', data_files=queries_path, cache_dir=self.cache_dir)['train']
|
||||
|
||||
queries = {e['id']: e['text'] for e in queries_data}
|
||||
return datasets.DatasetDict(queries)
|
||||
|
||||
def _download_file(self, download_url: str, save_dir: str):
|
||||
"""Download file from provided URL.
|
||||
|
||||
Args:
|
||||
download_url (str): Source URL of the file.
|
||||
save_dir (str): Path to the directory to save the zip file.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError
|
||||
|
||||
Returns:
|
||||
str: The path of the downloaded file.
|
||||
"""
|
||||
save_path = os.path.join(save_dir, download_url.split('/')[-1])
|
||||
|
||||
if self.force_redownload or (not os.path.exists(save_path) or os.path.getsize(save_path) == 0):
|
||||
cmd = ["wget", "-O", save_path, download_url]
|
||||
else:
|
||||
cmd = ["wget", "-nc", "-O", save_path, download_url]
|
||||
|
||||
try:
|
||||
subprocess.run(cmd, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning(e.output)
|
||||
|
||||
if not os.path.exists(save_path) or os.path.getsize(save_path) == 0:
|
||||
raise FileNotFoundError(f"Failed to download file from {download_url} to {save_path}")
|
||||
else:
|
||||
logger.info(f"Downloaded file from {download_url} to {save_path}")
|
||||
return save_path
|
||||
|
||||
def _get_fpath_size(self, fpath: str) -> int:
|
||||
"""Get the total size of the files in provided path.
|
||||
|
||||
Args:
|
||||
fpath (str): path of files to compute the size.
|
||||
|
||||
Returns:
|
||||
int: The total size in bytes.
|
||||
"""
|
||||
if not os.path.isdir(fpath):
|
||||
return os.path.getsize(fpath)
|
||||
else:
|
||||
total_size = 0
|
||||
for dirpath, _, filenames in os.walk(fpath):
|
||||
for f in filenames:
|
||||
fp = os.path.join(dirpath, f)
|
||||
total_size += os.path.getsize(fp)
|
||||
return total_size
|
||||
|
||||
def _download_gz_file(self, download_url: str, save_dir: str):
|
||||
"""Download and unzip the gzip file from provided URL.
|
||||
|
||||
Args:
|
||||
download_url (str): Source URL of the gzip file.
|
||||
save_dir (str): Path to the directory to save the gzip file.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError
|
||||
|
||||
Returns:
|
||||
str: The path to the file after unzip.
|
||||
"""
|
||||
gz_file_path = self._download_file(download_url, save_dir)
|
||||
cmd = ["gzip", "-d", gz_file_path]
|
||||
try:
|
||||
subprocess.run(cmd, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning(e.output)
|
||||
|
||||
file_path = gz_file_path.replace(".gz", "")
|
||||
if not os.path.exists(file_path) or self._get_fpath_size(file_path) == 0:
|
||||
raise FileNotFoundError(f"Failed to unzip file {gz_file_path}")
|
||||
|
||||
return file_path
|
||||
|
||||
def _download_zip_file(self, download_url: str, save_dir: str):
|
||||
"""Download and unzip the zip file from provided URL.
|
||||
|
||||
Args:
|
||||
download_url (str): Source URL of the zip file.
|
||||
save_dir (str): Path to the directory to save the zip file.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError
|
||||
|
||||
Returns:
|
||||
str: The path to the file after unzip.
|
||||
"""
|
||||
zip_file_path = self._download_file(download_url, save_dir)
|
||||
file_path = zip_file_path.replace(".zip", "")
|
||||
if self.force_redownload or not os.path.exists(file_path):
|
||||
cmd = ["unzip", "-o", zip_file_path, "-d", file_path]
|
||||
else:
|
||||
cmd = ["unzip", "-n", zip_file_path, "-d", file_path]
|
||||
|
||||
try:
|
||||
subprocess.run(cmd, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning(e.output)
|
||||
|
||||
if not os.path.exists(file_path) or self._get_fpath_size(file_path) == 0:
|
||||
raise FileNotFoundError(f"Failed to unzip file {zip_file_path}")
|
||||
|
||||
return file_path
|
||||
@@ -0,0 +1,500 @@
|
||||
"""
|
||||
Adapted from https://github.com/AIR-Bench/AIR-Bench/blob/0.1.0/air_benchmark/evaluation_utils/evaluator.py
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import json
|
||||
import pandas as pd
|
||||
from typing import Dict, Optional, List, Union
|
||||
|
||||
from .data_loader import AbsEvalDataLoader
|
||||
from .searcher import EvalRetriever, EvalReranker
|
||||
from .utils import evaluate_metrics, evaluate_mrr, evaluate_recall_cap
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsEvaluator:
|
||||
"""
|
||||
Base class of Evaluator.
|
||||
|
||||
Args:
|
||||
eval_name (str): The experiment name of current evaluation.
|
||||
data_loader (AbsEvalDataLoader): The data_loader to deal with data.
|
||||
overwrite (bool): If true, will overwrite the existing results.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
eval_name: str,
|
||||
data_loader: AbsEvalDataLoader,
|
||||
overwrite: bool = False,
|
||||
):
|
||||
self.eval_name = eval_name
|
||||
self.data_loader = data_loader
|
||||
self.overwrite = overwrite
|
||||
|
||||
def check_data_info(
|
||||
self,
|
||||
data_info: Dict[str, str],
|
||||
model_name: str,
|
||||
reranker_name: str,
|
||||
split: str,
|
||||
dataset_name: Optional[str] = None,
|
||||
):
|
||||
"""Check the validity of data info.
|
||||
|
||||
Args:
|
||||
data_info (Dict[str, str]): The loaded data info to be check.
|
||||
model_name (str): Name of model used.
|
||||
reranker_name (str): Name of reranker used.
|
||||
split (str): Split used in searching.
|
||||
dataset_name (Optional[str], optional): Name of dataset used. Defaults to None.
|
||||
|
||||
Raises:
|
||||
ValueError: eval_name mismatch
|
||||
ValueError: model_name or reranker_name mismatch
|
||||
ValueError: split mismatch
|
||||
ValueError: dataset_name mismatch
|
||||
"""
|
||||
if data_info["eval_name"] != self.eval_name:
|
||||
raise ValueError(
|
||||
f'eval_name mismatch: {data_info["eval_name"]} vs {self.eval_name}'
|
||||
)
|
||||
if (
|
||||
data_info["model_name"] != model_name
|
||||
or data_info["reranker_name"] != reranker_name
|
||||
):
|
||||
raise ValueError(
|
||||
f'model_name or reranker_name mismatch: {data_info["model_name"]} vs {model_name} or {data_info["reranker_name"]} vs {reranker_name}'
|
||||
)
|
||||
if (data_info["split"] != split):
|
||||
raise ValueError(
|
||||
f'split mismatch: {data_info["split"]} vs {split}'
|
||||
)
|
||||
if dataset_name is not None and data_info["dataset_name"] != dataset_name:
|
||||
raise ValueError(
|
||||
f'dataset_name mismatch: {data_info["dataset_name"]} vs {dataset_name}'
|
||||
)
|
||||
|
||||
def get_corpus_embd_save_dir(
|
||||
self,
|
||||
retriever_name: str,
|
||||
corpus_embd_save_dir: Optional[str] = None,
|
||||
dataset_name: Optional[str] = None
|
||||
):
|
||||
"""
|
||||
If corpus_embd_save_dir is not None, then it will be used as the base directory to save the corpus embeddings. For dataset such as MKQA,
|
||||
the corpus for all languages is the same, so the subclass can override this method to save the corpus embeddings in the same directory.
|
||||
|
||||
Args:
|
||||
retriever_name (str): Name of the retriever.
|
||||
corpus_embd_save_dir (str, optional): Directory that saving the corpus embedding.
|
||||
dataset_name (str, optional):
|
||||
"""
|
||||
if corpus_embd_save_dir is not None:
|
||||
if dataset_name is not None:
|
||||
corpus_embd_save_dir = os.path.join(corpus_embd_save_dir, retriever_name, dataset_name)
|
||||
else:
|
||||
corpus_embd_save_dir = os.path.join(corpus_embd_save_dir, retriever_name)
|
||||
return corpus_embd_save_dir
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
splits: Union[str, List[str]],
|
||||
search_results_save_dir: str,
|
||||
retriever: EvalRetriever,
|
||||
reranker: Optional[EvalReranker] = None,
|
||||
corpus_embd_save_dir: Optional[str] = None,
|
||||
ignore_identical_ids: bool = False,
|
||||
k_values: List[int] = [1, 3, 5, 10, 100, 1000],
|
||||
dataset_name: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""This is called during the evaluation process.
|
||||
|
||||
Args:
|
||||
splits (Union[str, List[str]]): Splits of datasets.
|
||||
search_results_save_dir (str): Directory to save the search results.
|
||||
retriever (EvalRetriever): object of :class:EvalRetriever.
|
||||
reranker (Optional[EvalReranker], optional): Object of :class:EvalReranker. Defaults to :data:`None`.
|
||||
corpus_embd_save_dir (Optional[str], optional): Directory to save the embedded corpus. Defaults to :data:`None`.
|
||||
ignore_identical_ids (bool, optional): If True, will ignore identical ids in search results. Defaults to :data:`False`.
|
||||
k_values (List[int], optional): Cutoffs. Defaults to :data:`[1, 3, 5, 10, 100, 1000]`.
|
||||
dataset_name (Optional[str], optional): Name of the datasets. Defaults to :data:`None`.
|
||||
"""
|
||||
# Check Splits
|
||||
checked_splits = self.data_loader.check_splits(splits, dataset_name=dataset_name)
|
||||
if len(checked_splits) == 0:
|
||||
logger.warning(f"{splits} not found in the dataset. Skipping evaluation.")
|
||||
return
|
||||
splits = checked_splits
|
||||
|
||||
if dataset_name is not None:
|
||||
save_name = f"{dataset_name}-" + "{split}.json"
|
||||
else:
|
||||
save_name = "{split}.json"
|
||||
|
||||
corpus_embd_save_dir = self.get_corpus_embd_save_dir(
|
||||
retriever_name=str(retriever),
|
||||
corpus_embd_save_dir=corpus_embd_save_dir,
|
||||
dataset_name=dataset_name
|
||||
)
|
||||
|
||||
# Retrieval Stage
|
||||
no_reranker_search_results_save_dir = os.path.join(
|
||||
search_results_save_dir, str(retriever), "NoReranker"
|
||||
)
|
||||
os.makedirs(no_reranker_search_results_save_dir, exist_ok=True)
|
||||
|
||||
flag = False
|
||||
for split in splits:
|
||||
split_no_reranker_search_results_save_path = os.path.join(
|
||||
no_reranker_search_results_save_dir, save_name.format(split=split)
|
||||
)
|
||||
if not os.path.exists(split_no_reranker_search_results_save_path) or self.overwrite:
|
||||
flag = True
|
||||
break
|
||||
|
||||
no_reranker_search_results_dict = {}
|
||||
if flag:
|
||||
corpus = self.data_loader.load_corpus(dataset_name=dataset_name)
|
||||
|
||||
queries_dict = {
|
||||
split: self.data_loader.load_queries(dataset_name=dataset_name, split=split)
|
||||
for split in splits
|
||||
}
|
||||
|
||||
all_queries = {}
|
||||
for _, split_queries in queries_dict.items():
|
||||
all_queries.update(split_queries)
|
||||
|
||||
all_no_reranker_search_results = retriever(
|
||||
corpus=corpus,
|
||||
queries=all_queries,
|
||||
corpus_embd_save_dir=corpus_embd_save_dir,
|
||||
ignore_identical_ids=ignore_identical_ids,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
for split in splits:
|
||||
split_queries = queries_dict[split]
|
||||
no_reranker_search_results_dict[split] = {
|
||||
qid: all_no_reranker_search_results[qid] for qid in split_queries
|
||||
}
|
||||
split_no_reranker_search_results_save_path = os.path.join(
|
||||
no_reranker_search_results_save_dir, save_name.format(split=split)
|
||||
)
|
||||
|
||||
self.save_search_results(
|
||||
eval_name=self.eval_name,
|
||||
model_name=str(retriever),
|
||||
reranker_name="NoReranker",
|
||||
search_results=no_reranker_search_results_dict[split],
|
||||
output_path=split_no_reranker_search_results_save_path,
|
||||
split=split,
|
||||
dataset_name=dataset_name,
|
||||
)
|
||||
else:
|
||||
for split in splits:
|
||||
split_no_reranker_search_results_save_path = os.path.join(
|
||||
no_reranker_search_results_save_dir, save_name.format(split=split)
|
||||
)
|
||||
data_info, search_results = self.load_search_results(split_no_reranker_search_results_save_path)
|
||||
|
||||
self.check_data_info(
|
||||
data_info=data_info,
|
||||
model_name=str(retriever),
|
||||
reranker_name="NoReranker",
|
||||
split=split,
|
||||
dataset_name=dataset_name,
|
||||
)
|
||||
no_reranker_search_results_dict[split] = search_results
|
||||
retriever.stop_multi_process_pool()
|
||||
eval_results_save_path = os.path.join(no_reranker_search_results_save_dir, 'EVAL', 'eval_results.json')
|
||||
if not os.path.exists(eval_results_save_path) or self.overwrite or flag:
|
||||
retriever_eval_results = self.evaluate_results(no_reranker_search_results_save_dir, k_values=k_values)
|
||||
self.output_eval_results_to_json(retriever_eval_results, eval_results_save_path)
|
||||
|
||||
# Reranking Stage
|
||||
if reranker is not None:
|
||||
reranker_search_results_save_dir = os.path.join(
|
||||
search_results_save_dir, str(retriever), str(reranker)
|
||||
)
|
||||
os.makedirs(reranker_search_results_save_dir, exist_ok=True)
|
||||
|
||||
corpus = self.data_loader.load_corpus(dataset_name=dataset_name)
|
||||
|
||||
queries_dict = {
|
||||
split: self.data_loader.load_queries(dataset_name=dataset_name, split=split)
|
||||
for split in splits
|
||||
}
|
||||
|
||||
flag = False
|
||||
for split in splits:
|
||||
rerank_search_results_save_path = os.path.join(
|
||||
reranker_search_results_save_dir, save_name.format(split=split)
|
||||
)
|
||||
|
||||
if os.path.exists(rerank_search_results_save_path) and not self.overwrite:
|
||||
continue
|
||||
|
||||
flag = True
|
||||
rerank_search_results = reranker(
|
||||
corpus=corpus,
|
||||
queries=queries_dict[split],
|
||||
search_results=no_reranker_search_results_dict[split],
|
||||
ignore_identical_ids=ignore_identical_ids,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
self.save_search_results(
|
||||
eval_name=self.eval_name,
|
||||
model_name=str(retriever),
|
||||
reranker_name=str(reranker),
|
||||
search_results=rerank_search_results,
|
||||
output_path=rerank_search_results_save_path,
|
||||
split=split,
|
||||
dataset_name=dataset_name,
|
||||
)
|
||||
reranker.stop_multi_process_pool()
|
||||
eval_results_save_path = os.path.join(reranker_search_results_save_dir, 'EVAL', 'eval_results.json')
|
||||
if not os.path.exists(eval_results_save_path) or self.overwrite or flag:
|
||||
reranker_eval_results = self.evaluate_results(reranker_search_results_save_dir, k_values=k_values)
|
||||
self.output_eval_results_to_json(reranker_eval_results, eval_results_save_path)
|
||||
|
||||
@staticmethod
|
||||
def save_search_results(
|
||||
eval_name: str,
|
||||
model_name: str,
|
||||
reranker_name: str,
|
||||
search_results: Dict[str, Dict[str, float]],
|
||||
output_path: str,
|
||||
split: str,
|
||||
dataset_name: Optional[str] = None,
|
||||
):
|
||||
"""Save the metadata and search results into a file.
|
||||
|
||||
Args:
|
||||
eval_name (str): The experiment name of current evaluation.
|
||||
model_name (str): Name of model used.
|
||||
reranker_name (str): Name of reranker used.
|
||||
search_results (Dict[str, Dict[str, float]]): Dictionary of search results.
|
||||
output_path (str): Output path to write the results.
|
||||
split (str): Split used in searching.
|
||||
dataset_name (Optional[str], optional): Name of dataset used. Defaults to :data:`None`.
|
||||
"""
|
||||
data = {
|
||||
"eval_name": eval_name,
|
||||
"model_name": model_name,
|
||||
"reranker_name": reranker_name,
|
||||
"split": split,
|
||||
"dataset_name": dataset_name,
|
||||
"search_results": search_results,
|
||||
}
|
||||
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
with open(output_path, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=4)
|
||||
|
||||
@staticmethod
|
||||
def load_search_results(input_path: str):
|
||||
"""Load search results from path.
|
||||
|
||||
Args:
|
||||
input_path (str): Path to load from.
|
||||
|
||||
Returns:
|
||||
dict, dict: data info that contains metadata and search results.
|
||||
"""
|
||||
with open(input_path, "r", encoding="utf-8") as f:
|
||||
data_info = json.load(f)
|
||||
|
||||
search_results = data_info.pop("search_results")
|
||||
return data_info, search_results
|
||||
|
||||
@staticmethod
|
||||
def compute_metrics(
|
||||
qrels: Dict[str, Dict[str, int]],
|
||||
search_results: Dict[str, Dict[str, float]],
|
||||
k_values: List[int],
|
||||
):
|
||||
"""Evaluate the model with metrics.
|
||||
|
||||
Args:
|
||||
qrels (Dict[str, Dict[str, int]]): Ground truth relevance of queries and documents.
|
||||
search_results (Dict[str, Dict[str, float]]): Dictionary of search results
|
||||
k_values (List[int]): Cutoffs.
|
||||
|
||||
Returns:
|
||||
dict: The results of the metrics.
|
||||
"""
|
||||
ndcg, _map, recall, precision = evaluate_metrics(
|
||||
qrels=qrels,
|
||||
results=search_results,
|
||||
k_values=k_values,
|
||||
)
|
||||
mrr = evaluate_mrr(
|
||||
qrels=qrels,
|
||||
results=search_results,
|
||||
k_values=k_values,
|
||||
)
|
||||
recall_cap = evaluate_recall_cap(
|
||||
qrels=qrels,
|
||||
results=search_results,
|
||||
k_values=k_values,
|
||||
)
|
||||
scores = {
|
||||
**{f"ndcg_at_{k.split('@')[1]}": v for (k, v) in ndcg.items()},
|
||||
**{f"map_at_{k.split('@')[1]}": v for (k, v) in _map.items()},
|
||||
**{f"recall_at_{k.split('@')[1]}": v for (k, v) in recall.items()},
|
||||
**{f"precision_at_{k.split('@')[1]}": v for (k, v) in precision.items()},
|
||||
**{f"mrr_at_{k.split('@')[1]}": v for (k, v) in mrr.items()},
|
||||
**{f"recall_cap_at_{k.split('@')[1]}": v for (k, v) in recall_cap.items()},
|
||||
}
|
||||
return scores
|
||||
|
||||
def evaluate_results(
|
||||
self,
|
||||
search_results_save_dir: str,
|
||||
k_values: List[int] = [1, 3, 5, 10, 100, 1000]
|
||||
):
|
||||
"""Compute metrics according to the results in the directory.
|
||||
|
||||
Args:
|
||||
search_results_save_dir (str): Path to the search results.
|
||||
k_values (List[int], optional): Cutoffs. Defaults to :data:`[1, 3, 5, 10, 100, 1000]`.
|
||||
|
||||
Returns:
|
||||
dict: Evaluation results.
|
||||
"""
|
||||
eval_results_dict = {}
|
||||
|
||||
for file in os.listdir(search_results_save_dir):
|
||||
if not file.endswith('.json'):
|
||||
continue
|
||||
|
||||
file_path = os.path.join(search_results_save_dir, file)
|
||||
data_info, search_results = self.load_search_results(file_path)
|
||||
|
||||
_eval_name = data_info['eval_name']
|
||||
assert _eval_name == self.eval_name, f'Mismatch eval_name: {_eval_name} vs {self.eval_name} in {file_path}'
|
||||
|
||||
split = data_info['split']
|
||||
dataset_name = data_info.get('dataset_name', None)
|
||||
qrels = self.data_loader.load_qrels(dataset_name=dataset_name, split=split)
|
||||
|
||||
eval_results = self.compute_metrics(
|
||||
qrels=qrels,
|
||||
search_results=search_results,
|
||||
k_values=k_values
|
||||
)
|
||||
|
||||
if dataset_name is not None:
|
||||
key = f"{dataset_name}-{split}"
|
||||
else:
|
||||
key = split
|
||||
eval_results_dict[key] = eval_results
|
||||
|
||||
return eval_results_dict
|
||||
|
||||
@staticmethod
|
||||
def output_eval_results_to_json(eval_results_dict: dict, output_path: str):
|
||||
"""Write the evaluation results into a json file.
|
||||
|
||||
Args:
|
||||
eval_results_dict (dict): Dictionary of the evaluation results.
|
||||
output_path (str): Output path to write the json file.
|
||||
"""
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(eval_results_dict, f, indent=4)
|
||||
logger.info(f"Results saved to {output_path}")
|
||||
|
||||
@staticmethod
|
||||
def get_results_df(metric: str, eval_results_dict: dict):
|
||||
"""Get the results from dictionary to a DataFrame.
|
||||
|
||||
Args:
|
||||
metric (str): Selected metric.
|
||||
eval_results_dict (dict): Dictionary of the evaluation results.
|
||||
|
||||
Returns:
|
||||
DataFrame: DataFrame of the results.
|
||||
"""
|
||||
results_dict = {}
|
||||
|
||||
for model_name, model_results in eval_results_dict.items():
|
||||
results_dict[model_name] = {}
|
||||
for reranker_name, reranker_results in model_results.items():
|
||||
results_dict[model_name][reranker_name] = {}
|
||||
for split, split_results in reranker_results.items():
|
||||
if metric in split_results:
|
||||
results_dict[model_name][reranker_name][split] = split_results[metric]
|
||||
else:
|
||||
results_dict[model_name][reranker_name][split] = None
|
||||
|
||||
model_reranker_pairs = set()
|
||||
all_splits = set()
|
||||
for model_name, model_results in results_dict.items():
|
||||
for reranker_name, reranker_results in model_results.items():
|
||||
model_reranker_pairs.add((model_name, reranker_name))
|
||||
all_splits.update(reranker_results.keys())
|
||||
|
||||
index = [(model, reranker) for model, reranker in model_reranker_pairs]
|
||||
multi_index = pd.MultiIndex.from_tuples(index, names=['Model', 'Reranker'])
|
||||
|
||||
all_splits = sorted(list(all_splits))
|
||||
overall_columns = ['average'] + all_splits
|
||||
overall_df = pd.DataFrame(index=multi_index, columns=overall_columns)
|
||||
|
||||
for model, reranker in model_reranker_pairs:
|
||||
for split in all_splits:
|
||||
if model in results_dict and reranker in results_dict[model] and split in results_dict[model][reranker]:
|
||||
overall_df.loc[(model, reranker), split] = results_dict[model][reranker][split]
|
||||
else:
|
||||
overall_df.loc[(model, reranker), split] = None
|
||||
if overall_df.loc[(model, reranker), all_splits].isnull().any():
|
||||
overall_df.loc[(model, reranker), 'average'] = None
|
||||
else:
|
||||
overall_df.loc[(model, reranker), 'average'] = overall_df.loc[(model, reranker), all_splits].mean()
|
||||
|
||||
return overall_df
|
||||
|
||||
@staticmethod
|
||||
def output_eval_results_to_markdown(eval_results_dict: dict, output_path: str, metrics: Union[List[str], str]):
|
||||
"""Write the evaluation results to a markdown file.
|
||||
|
||||
Args:
|
||||
eval_results_dict (dict): Dictionary that contains evaluation results.
|
||||
output_path (str): Path to write the output to.
|
||||
metrics (Union[List[str], str]): The metrics that will be written in the markdown file.
|
||||
"""
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
if isinstance(metrics, str):
|
||||
metrics = [metrics]
|
||||
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
for metric in metrics:
|
||||
f.write(f"## {metric}\n\n")
|
||||
results_df = AbsEvaluator.get_results_df(metric, eval_results_dict)
|
||||
max_index = dict(results_df.idxmax(axis=0))
|
||||
splits = results_df.columns
|
||||
f.write(f"| Model | Reranker | {' | '.join(splits)} |\n")
|
||||
f.write(f"| :---- | :---- | {' | '.join([':---:' for _ in splits])} |\n")
|
||||
for i, row in results_df.iterrows():
|
||||
line = f"| {i[0]} | {i[1]} | "
|
||||
for s, v in row.items():
|
||||
if v is None:
|
||||
line += "- | "
|
||||
else:
|
||||
if i != max_index[s]:
|
||||
line += f'{v*100:.3f} | '
|
||||
else:
|
||||
line += f'**{v*100:.3f}** | '
|
||||
f.write(line + "\n")
|
||||
f.write("\n")
|
||||
logger.info(f"Results saved to {output_path}")
|
||||
@@ -0,0 +1,229 @@
|
||||
import os
|
||||
import json
|
||||
import logging
|
||||
from typing import List, Union, Tuple
|
||||
|
||||
from FlagEmbedding import FlagAutoModel, FlagAutoReranker, AbsEmbedder, AbsReranker
|
||||
|
||||
from .arguments import AbsEvalArgs, AbsEvalModelArgs
|
||||
from .evaluator import AbsEvaluator
|
||||
from .searcher import EvalDenseRetriever, EvalReranker
|
||||
from .data_loader import AbsEvalDataLoader
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsEvalRunner:
|
||||
"""
|
||||
Abstract class of evaluation runner.
|
||||
|
||||
Args:
|
||||
eval_args (AbsEvalArgs): :class:AbsEvalArgs object with the evaluation arguments.
|
||||
model_args (AbsEvalModelArgs): :class:AbsEvalModelArgs object with the model arguments.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
eval_args: AbsEvalArgs,
|
||||
model_args: AbsEvalModelArgs,
|
||||
):
|
||||
self.eval_args = eval_args
|
||||
self.model_args = model_args
|
||||
|
||||
self.retriever, self.reranker = self.load_retriever_and_reranker()
|
||||
self.data_loader = self.load_data_loader()
|
||||
self.evaluator = self.load_evaluator()
|
||||
|
||||
@staticmethod
|
||||
def get_models(model_args: AbsEvalModelArgs) -> Tuple[AbsEmbedder, Union[AbsReranker, None]]:
|
||||
"""Get the embedding and reranker model
|
||||
|
||||
Args:
|
||||
model_args (AbsEvalModelArgs): :class:AbsEvalModelArgs object with the model arguments.
|
||||
|
||||
Returns:
|
||||
Tuple[AbsEmbedder, Union[AbsReranker, None]]: A :class:AbsEmbedder object of embedding model, and
|
||||
:class:AbsReranker object of reranker model if path provided.
|
||||
"""
|
||||
embedder = FlagAutoModel.from_finetuned(
|
||||
model_name_or_path=model_args.embedder_name_or_path,
|
||||
model_class=model_args.embedder_model_class,
|
||||
normalize_embeddings=model_args.normalize_embeddings,
|
||||
pooling_method=model_args.pooling_method,
|
||||
use_fp16=model_args.use_fp16,
|
||||
use_bf16=model_args.use_bf16,
|
||||
query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
|
||||
query_instruction_format=model_args.query_instruction_format_for_retrieval,
|
||||
devices=model_args.devices,
|
||||
examples_for_task=model_args.examples_for_task,
|
||||
examples_instruction_format=model_args.examples_instruction_format,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
cache_dir=model_args.cache_dir,
|
||||
domain_for_pseudo_moe=model_args.domain_for_pseudo_moe,
|
||||
batch_size=model_args.embedder_batch_size,
|
||||
query_max_length=model_args.embedder_query_max_length,
|
||||
passage_max_length=model_args.embedder_passage_max_length,
|
||||
truncate_dim=model_args.truncate_dim,
|
||||
)
|
||||
embedder.model.config._name_or_path = model_args.embedder_name_or_path
|
||||
reranker = None
|
||||
if model_args.reranker_name_or_path is not None:
|
||||
reranker = FlagAutoReranker.from_finetuned(
|
||||
model_name_or_path=model_args.reranker_name_or_path,
|
||||
model_class=model_args.reranker_model_class,
|
||||
peft_path=model_args.reranker_peft_path,
|
||||
use_fp16=model_args.use_fp16,
|
||||
use_bf16=model_args.use_bf16,
|
||||
query_instruction_for_rerank=model_args.query_instruction_for_rerank,
|
||||
query_instruction_format=model_args.query_instruction_format_for_rerank,
|
||||
passage_instruction_for_rerank=model_args.passage_instruction_for_rerank,
|
||||
passage_instruction_format=model_args.passage_instruction_format_for_rerank,
|
||||
cache_dir=model_args.cache_dir,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
devices=model_args.devices,
|
||||
normalize=model_args.normalize,
|
||||
prompt=model_args.prompt,
|
||||
cutoff_layers=model_args.cutoff_layers,
|
||||
compress_layers=model_args.compress_layers,
|
||||
compress_ratio=model_args.compress_ratio,
|
||||
batch_size=model_args.reranker_batch_size,
|
||||
query_max_length=model_args.reranker_query_max_length,
|
||||
max_length=model_args.reranker_max_length,
|
||||
)
|
||||
reranker.model.config._name_or_path = model_args.reranker_name_or_path
|
||||
return embedder, reranker
|
||||
|
||||
def load_retriever_and_reranker(self) -> Tuple[EvalDenseRetriever, Union[EvalReranker, None]]:
|
||||
"""Load retriever and reranker for evaluation
|
||||
|
||||
Returns:
|
||||
Tuple[EvalDenseRetriever, Union[EvalReranker, None]]: A :class:EvalDenseRetriever object for retrieval, and a
|
||||
:class:EvalReranker object if reranker provided.
|
||||
"""
|
||||
embedder, reranker = self.get_models(self.model_args)
|
||||
retriever = EvalDenseRetriever(
|
||||
embedder,
|
||||
search_top_k=self.eval_args.search_top_k,
|
||||
overwrite=self.eval_args.overwrite
|
||||
)
|
||||
if reranker is not None:
|
||||
reranker = EvalReranker(reranker, rerank_top_k=self.eval_args.rerank_top_k)
|
||||
return retriever, reranker
|
||||
|
||||
def load_data_loader(self) -> AbsEvalDataLoader:
|
||||
"""Load the data loader
|
||||
|
||||
Returns:
|
||||
AbsEvalDataLoader: Data loader object for that specific task.
|
||||
"""
|
||||
data_loader = AbsEvalDataLoader(
|
||||
eval_name=self.eval_args.eval_name,
|
||||
dataset_dir=self.eval_args.dataset_dir,
|
||||
cache_dir=self.eval_args.cache_path,
|
||||
token=self.eval_args.token,
|
||||
force_redownload=self.eval_args.force_redownload,
|
||||
)
|
||||
return data_loader
|
||||
|
||||
def load_evaluator(self) -> AbsEvaluator:
|
||||
"""Load the evaluator for evaluation
|
||||
|
||||
Returns:
|
||||
AbsEvaluator: the evaluator to run the evaluation.
|
||||
"""
|
||||
evaluator = AbsEvaluator(
|
||||
eval_name=self.eval_args.eval_name,
|
||||
data_loader=self.data_loader,
|
||||
overwrite=self.eval_args.overwrite,
|
||||
)
|
||||
return evaluator
|
||||
|
||||
@staticmethod
|
||||
def evaluate_metrics(
|
||||
search_results_save_dir: str,
|
||||
output_method: str = "markdown",
|
||||
output_path: str = "./eval_dev_results.md",
|
||||
metrics: Union[str, List[str]] = ["ndcg_at_10", "recall_at_10"]
|
||||
):
|
||||
"""Evaluate the provided metrics and write the results.
|
||||
|
||||
Args:
|
||||
search_results_save_dir (str): Path to save the search results.
|
||||
output_method (str, optional): Output results to `json` or `markdown`. Defaults to :data:`"markdown"`.
|
||||
output_path (str, optional): Path to write the output. Defaults to :data:`"./eval_dev_results.md"`.
|
||||
metrics (Union[str, List[str]], optional): metrics to use. Defaults to :data:`["ndcg_at_10", "recall_at_10"]`.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: Eval results not found
|
||||
ValueError: Invalid output method
|
||||
"""
|
||||
eval_results_dict = {}
|
||||
for model_name in sorted(os.listdir(search_results_save_dir)):
|
||||
model_search_results_save_dir = os.path.join(search_results_save_dir, model_name)
|
||||
if not os.path.isdir(model_search_results_save_dir):
|
||||
continue
|
||||
for reranker_name in sorted(os.listdir(model_search_results_save_dir)):
|
||||
reranker_search_results_save_dir = os.path.join(model_search_results_save_dir, reranker_name)
|
||||
if not os.path.isdir(reranker_search_results_save_dir):
|
||||
continue
|
||||
eval_results_path = os.path.join(reranker_search_results_save_dir, 'EVAL', "eval_results.json")
|
||||
if os.path.exists(eval_results_path):
|
||||
eval_results = json.load(open(eval_results_path, encoding='utf-8'))
|
||||
else:
|
||||
logger.warning(f"Eval results not found: {eval_results_path}")
|
||||
continue
|
||||
|
||||
if model_name not in eval_results_dict:
|
||||
eval_results_dict[model_name] = {}
|
||||
eval_results_dict[model_name][reranker_name] = eval_results
|
||||
|
||||
if output_method == "json":
|
||||
AbsEvaluator.output_eval_results_to_json(eval_results_dict, output_path)
|
||||
elif output_method == "markdown":
|
||||
AbsEvaluator.output_eval_results_to_markdown(eval_results_dict, output_path, metrics)
|
||||
else:
|
||||
raise ValueError(f"Invalid output method: {output_method}. Available methods: ['json', 'markdown']")
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
Run the whole evaluation.
|
||||
"""
|
||||
if self.eval_args.dataset_names is None:
|
||||
dataset_names = self.data_loader.available_dataset_names()
|
||||
else:
|
||||
dataset_names = self.data_loader.check_dataset_names(self.eval_args.dataset_names)
|
||||
|
||||
if len(dataset_names) == 0:
|
||||
logger.info(f"Running {self.eval_args.eval_name} evaluation on the default dataset.")
|
||||
self.evaluator(
|
||||
splits=self.eval_args.splits,
|
||||
search_results_save_dir=self.eval_args.output_dir,
|
||||
retriever=self.retriever,
|
||||
reranker=self.reranker,
|
||||
corpus_embd_save_dir=self.eval_args.corpus_embd_save_dir,
|
||||
ignore_identical_ids=self.eval_args.ignore_identical_ids,
|
||||
k_values=self.eval_args.k_values
|
||||
)
|
||||
logger.info(f"{self.eval_args.eval_name} evaluation completed.")
|
||||
else:
|
||||
logger.info(f"Running {self.eval_args.eval_name} evaluation on the following dataset names: {dataset_names}")
|
||||
for dataset_name in dataset_names:
|
||||
logger.info(f"Running {self.eval_args.eval_name} evaluation on: {dataset_name}")
|
||||
self.evaluator(
|
||||
splits=self.eval_args.splits,
|
||||
search_results_save_dir=self.eval_args.output_dir,
|
||||
retriever=self.retriever,
|
||||
reranker=self.reranker,
|
||||
corpus_embd_save_dir=self.eval_args.corpus_embd_save_dir,
|
||||
ignore_identical_ids=self.eval_args.ignore_identical_ids,
|
||||
k_values=self.eval_args.k_values,
|
||||
dataset_name=dataset_name,
|
||||
)
|
||||
logger.info(f"{self.eval_args.eval_name} evaluation on {dataset_names} completed.")
|
||||
|
||||
logger.info("Start computing metrics.")
|
||||
self.evaluate_metrics(
|
||||
search_results_save_dir=self.eval_args.output_dir,
|
||||
output_method=self.eval_args.eval_output_method,
|
||||
output_path=self.eval_args.eval_output_path,
|
||||
metrics=self.eval_args.eval_metrics
|
||||
)
|
||||
@@ -0,0 +1,248 @@
|
||||
"""
|
||||
Adapted from https://github.com/AIR-Bench/AIR-Bench/blob/0.1.0/air_benchmark/evaluation_utils/searcher.py
|
||||
"""
|
||||
import os
|
||||
import logging
|
||||
import gc
|
||||
import torch
|
||||
import numpy as np
|
||||
from typing import Any, Dict, Optional
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from FlagEmbedding.abc.inference import AbsEmbedder, AbsReranker
|
||||
from FlagEmbedding.abc.evaluation.utils import index, search
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class EvalRetriever(ABC):
|
||||
"""
|
||||
This is the base class for retriever.
|
||||
"""
|
||||
def __init__(self, embedder: AbsEmbedder, search_top_k: int = 1000, overwrite: bool = False):
|
||||
self.embedder = embedder
|
||||
self.search_top_k = search_top_k
|
||||
self.overwrite = overwrite
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""
|
||||
Returns: str: Name of the retriever.
|
||||
"""
|
||||
return os.path.basename(self.embedder.model.config._name_or_path)
|
||||
|
||||
def stop_multi_process_pool(self):
|
||||
self.embedder.stop_self_pool()
|
||||
# if self.embedder.pool is not None:
|
||||
# self.embedder.stop_multi_process_pool(self.embedder.pool)
|
||||
# self.embedder.pool = None
|
||||
# self.embedder.model.to('cpu')
|
||||
# gc.collect()
|
||||
# torch.cuda.empty_cache()
|
||||
|
||||
@abstractmethod
|
||||
def __call__(
|
||||
self,
|
||||
corpus: Dict[str, Dict[str, Any]],
|
||||
queries: Dict[str, str],
|
||||
corpus_embd_save_dir: Optional[str] = None,
|
||||
ignore_identical_ids: bool = False,
|
||||
**kwargs,
|
||||
) -> Dict[str, Dict[str, float]]:
|
||||
"""
|
||||
Abstract method to be overrode. This is called during the retrieval process.
|
||||
|
||||
Parameters:
|
||||
corpus: Dict[str, Dict[str, Any]]: Corpus of documents.
|
||||
Structure: {<docid>: {"text": <text>}}.
|
||||
Example: {"doc-0": {"text": "This is a document."}}
|
||||
queries: Dict[str, str]: Queries to search for.
|
||||
Structure: {<qid>: <query>}.
|
||||
Example: {"q-0": "This is a query."}
|
||||
corpus_embd_save_dir (Optional[str]): Defaults to :data:`None`.
|
||||
ignore_identical_ids (bool): Defaults to :data:`False`.
|
||||
**kwargs: Any: Additional arguments.
|
||||
|
||||
Returns: Dict[str, Dict[str, float]]: Top-k search results for each query. k is specified by search_top_k.
|
||||
Structure: {qid: {docid: score}}. The higher is the score, the more relevant is the document.
|
||||
Example: {"q-0": {"doc-0": 0.9}}
|
||||
"""
|
||||
|
||||
|
||||
class EvalDenseRetriever(EvalRetriever):
|
||||
"""
|
||||
Child class of :class:EvalRetriever for dense retrieval.
|
||||
"""
|
||||
def __call__(
|
||||
self,
|
||||
corpus: Dict[str, Dict[str, Any]],
|
||||
queries: Dict[str, str],
|
||||
corpus_embd_save_dir: Optional[str] = None,
|
||||
ignore_identical_ids: bool = False,
|
||||
**kwargs,
|
||||
) -> Dict[str, Dict[str, float]]:
|
||||
"""
|
||||
This is called during the retrieval process.
|
||||
|
||||
Parameters:
|
||||
corpus: Dict[str, Dict[str, Any]]: Corpus of documents.
|
||||
Structure: {<docid>: {"text": <text>}}.
|
||||
Example: {"doc-0": {"text": "This is a document."}}
|
||||
queries: Dict[str, str]: Queries to search for.
|
||||
Structure: {<qid>: <query>}.
|
||||
Example: {"q-0": "This is a query."}
|
||||
corpus_embd_save_dir (Optional[str]): Defaults to :data:`None`.
|
||||
ignore_identical_ids (bool): Defaults to :data:`False`.
|
||||
**kwargs: Any: Additional arguments.
|
||||
|
||||
Returns: Dict[str, Dict[str, float]]: Top-k search results for each query. k is specified by search_top_k.
|
||||
Structure: {qid: {docid: score}}. The higher is the score, the more relevant is the document.
|
||||
Example: {"q-0": {"doc-0": 0.9}}
|
||||
"""
|
||||
if ignore_identical_ids:
|
||||
logger.warning("ignore_identical_ids is set to True. This means that the search results will not contain identical ids. Note: Dataset such as MIRACL should NOT set this to True.")
|
||||
|
||||
# dense embedding models do not require language as input: AIRBench evaluation
|
||||
kwargs.pop("language", None)
|
||||
|
||||
corpus_ids = []
|
||||
corpus_texts = []
|
||||
for docid, doc in corpus.items():
|
||||
corpus_ids.append(docid)
|
||||
corpus_texts.append(
|
||||
doc["text"] if "title" not in doc
|
||||
else f"{doc['title']} {doc['text']}".strip()
|
||||
)
|
||||
queries_ids = []
|
||||
queries_texts = []
|
||||
for qid, query in queries.items():
|
||||
queries_ids.append(qid)
|
||||
queries_texts.append(query)
|
||||
|
||||
if corpus_embd_save_dir is not None:
|
||||
if os.path.exists(os.path.join(corpus_embd_save_dir, "doc.npy")) and not self.overwrite:
|
||||
corpus_emb = np.load(os.path.join(corpus_embd_save_dir, "doc.npy"))
|
||||
else:
|
||||
corpus_emb = self.embedder.encode_corpus(corpus_texts, **kwargs)
|
||||
else:
|
||||
corpus_emb = self.embedder.encode_corpus(corpus_texts, **kwargs)
|
||||
|
||||
queries_emb = self.embedder.encode_queries(queries_texts, **kwargs)
|
||||
|
||||
# check if the embeddings are in dictionary format: M3Embedder
|
||||
if isinstance(corpus_emb, dict):
|
||||
corpus_emb = corpus_emb["dense_vecs"]
|
||||
if isinstance(queries_emb, dict):
|
||||
queries_emb = queries_emb["dense_vecs"]
|
||||
|
||||
if corpus_embd_save_dir is not None and \
|
||||
(not os.path.exists(os.path.join(corpus_embd_save_dir, "doc.npy")) or self.overwrite):
|
||||
os.makedirs(corpus_embd_save_dir, exist_ok=True)
|
||||
np.save(os.path.join(corpus_embd_save_dir, "doc.npy"), corpus_emb)
|
||||
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
faiss_index = index(corpus_embeddings=corpus_emb)
|
||||
all_scores, all_indices = search(query_embeddings=queries_emb, faiss_index=faiss_index, k=self.search_top_k)
|
||||
|
||||
results = {}
|
||||
for idx, (scores, indices) in enumerate(zip(all_scores, all_indices)):
|
||||
results[queries_ids[idx]] = {}
|
||||
for score, indice in zip(scores, indices):
|
||||
if indice != -1:
|
||||
if ignore_identical_ids and corpus_ids[indice] == queries_ids[idx]:
|
||||
continue
|
||||
results[queries_ids[idx]][corpus_ids[indice]] = float(score)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
class EvalReranker:
|
||||
"""
|
||||
Class for reranker during evaluation.
|
||||
"""
|
||||
def __init__(self, reranker: AbsReranker, rerank_top_k: int = 100):
|
||||
self.reranker = reranker
|
||||
self.rerank_top_k = rerank_top_k
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""
|
||||
Returns: str: Name of the reranker.
|
||||
"""
|
||||
return os.path.basename(self.reranker.model.config._name_or_path)
|
||||
|
||||
def stop_multi_process_pool(self):
|
||||
self.reranker.stop_self_pool()
|
||||
# if self.reranker.pool is not None:
|
||||
# self.reranker.stop_multi_process_pool(self.reranker.pool)
|
||||
# self.reranker.pool = None
|
||||
# self.reranker.model.to('cpu')
|
||||
# gc.collect()
|
||||
# torch.cuda.empty_cache()
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
corpus: Dict[str, Dict[str, Any]],
|
||||
queries: Dict[str, str],
|
||||
search_results: Dict[str, Dict[str, float]],
|
||||
ignore_identical_ids: bool = False,
|
||||
**kwargs,
|
||||
) -> Dict[str, Dict[str, float]]:
|
||||
"""
|
||||
This is called during the reranking process.
|
||||
|
||||
Parameters:
|
||||
corpus: Dict[str, Dict[str, Any]]: Corpus of documents.
|
||||
Structure: {<docid>: {"text": <text>}}.
|
||||
Example: {"doc-0": {"text": "This is a document."}}
|
||||
queries: Dict[str, str]: Queries to search for.
|
||||
Structure: {<qid>: <query>}.
|
||||
Example: {"q-0": "This is a query."}
|
||||
search_results: Dict[str, Dict[str, float]]: Search results for each query.
|
||||
Structure: {qid: {docid: score}}. The higher is the score, the more relevant is the document.
|
||||
Example: {"q-0": {"doc-0": 0.9}}
|
||||
**kwargs: Any: Additional arguments.
|
||||
|
||||
Returns: Dict[str, Dict[str, float]]: Reranked search results for each query. k is specified by rerank_top_k.
|
||||
Structure: {qid: {docid: score}}. The higher is the score, the more relevant is the document.
|
||||
Example: {"q-0": {"doc-0": 0.9}}
|
||||
"""
|
||||
# truncate search results to top_k
|
||||
for qid in search_results:
|
||||
search_results[qid] = dict(
|
||||
sorted(search_results[qid].items(), key=lambda x: x[1], reverse=True)[
|
||||
:self.rerank_top_k
|
||||
]
|
||||
)
|
||||
# generate sentence pairs
|
||||
sentence_pairs = []
|
||||
pairs = []
|
||||
for qid in search_results:
|
||||
for docid in search_results[qid]:
|
||||
if ignore_identical_ids and qid == docid:
|
||||
continue
|
||||
sentence_pairs.append(
|
||||
{
|
||||
"qid": qid,
|
||||
"docid": docid,
|
||||
"query": queries[qid],
|
||||
"doc": corpus[docid]["text"] if "title" not in corpus[docid]
|
||||
else f"{corpus[docid]['title']} {corpus[docid]['text']}".strip(),
|
||||
}
|
||||
)
|
||||
pairs.append(
|
||||
(
|
||||
queries[qid],
|
||||
corpus[docid]["text"] if "title" not in corpus[docid]
|
||||
else f"{corpus[docid]['title']} {corpus[docid]['text']}".strip()
|
||||
)
|
||||
)
|
||||
# compute scores
|
||||
scores = self.reranker.compute_score(pairs)
|
||||
for i, score in enumerate(scores):
|
||||
sentence_pairs[i]["score"] = float(score)
|
||||
# rerank
|
||||
reranked_results = {qid: {} for qid in search_results}
|
||||
for pair in sentence_pairs:
|
||||
reranked_results[pair["qid"]][pair["docid"]] = pair["score"]
|
||||
return reranked_results
|
||||
@@ -0,0 +1,228 @@
|
||||
import faiss
|
||||
import torch
|
||||
import logging
|
||||
import numpy as np
|
||||
import pytrec_eval
|
||||
from tqdm import tqdm
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Tuple, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Modified from https://github.com/beir-cellar/beir/blob/f062f038c4bfd19a8ca942a9910b1e0d218759d4/beir/retrieval/custom_metrics.py#L4
|
||||
def evaluate_mrr(
|
||||
qrels: Dict[str, Dict[str, int]],
|
||||
results: Dict[str, Dict[str, float]],
|
||||
k_values: List[int],
|
||||
) -> Tuple[Dict[str, float]]:
|
||||
"""Compute mean reciprocal rank (MRR).
|
||||
|
||||
Args:
|
||||
qrels (Dict[str, Dict[str, int]]): Ground truth relevance.
|
||||
results (Dict[str, Dict[str, float]]): Search results to evaluate.
|
||||
k_values (List[int]): Cutoffs.
|
||||
|
||||
Returns:
|
||||
Tuple[Dict[str, float]]: MRR results at provided k values.
|
||||
"""
|
||||
mrr = defaultdict(list)
|
||||
|
||||
k_max, top_hits = max(k_values), {}
|
||||
|
||||
for query_id, doc_scores in results.items():
|
||||
top_hits[query_id] = sorted(
|
||||
doc_scores.items(), key=lambda item: item[1], reverse=True
|
||||
)[0:k_max]
|
||||
|
||||
for query_id in top_hits:
|
||||
query_relevant_docs = {
|
||||
doc_id for doc_id in qrels[query_id] if qrels[query_id][doc_id] > 0
|
||||
}
|
||||
for k in k_values:
|
||||
rr = 0
|
||||
for rank, hit in enumerate(top_hits[query_id][0:k], 1):
|
||||
if hit[0] in query_relevant_docs:
|
||||
rr = 1.0 / rank
|
||||
break
|
||||
mrr[f"MRR@{k}"].append(rr)
|
||||
|
||||
for k in k_values:
|
||||
mrr[f"MRR@{k}"] = round(sum(mrr[f"MRR@{k}"]) / len(qrels), 5)
|
||||
return mrr
|
||||
|
||||
|
||||
# Modified from https://github.com/beir-cellar/beir/blob/f062f038c4bfd19a8ca942a9910b1e0d218759d4/beir/retrieval/custom_metrics.py#L33
|
||||
def evaluate_recall_cap(
|
||||
qrels: Dict[str, Dict[str, int]],
|
||||
results: Dict[str, Dict[str, float]],
|
||||
k_values: List[int]
|
||||
) -> Tuple[Dict[str, float]]:
|
||||
"""Compute capped recall.
|
||||
|
||||
Args:
|
||||
qrels (Dict[str, Dict[str, int]]): Ground truth relevance.
|
||||
results (Dict[str, Dict[str, float]]): Search results to evaluate.
|
||||
k_values (List[int]): Cutoffs.
|
||||
|
||||
Returns:
|
||||
Tuple[Dict[str, float]]: Capped recall results at provided k values.
|
||||
"""
|
||||
capped_recall = {}
|
||||
|
||||
for k in k_values:
|
||||
capped_recall[f"R_cap@{k}"] = 0.0
|
||||
|
||||
k_max = max(k_values)
|
||||
logging.info("\n")
|
||||
|
||||
for query_id, doc_scores in results.items():
|
||||
top_hits = sorted(doc_scores.items(), key=lambda item: item[1], reverse=True)[0:k_max]
|
||||
query_relevant_docs = [doc_id for doc_id in qrels[query_id] if qrels[query_id][doc_id] > 0]
|
||||
for k in k_values:
|
||||
retrieved_docs = [row[0] for row in top_hits[0:k] if qrels[query_id].get(row[0], 0) > 0]
|
||||
denominator = min(len(query_relevant_docs), k)
|
||||
capped_recall[f"R_cap@{k}"] += (len(retrieved_docs) / denominator)
|
||||
|
||||
for k in k_values:
|
||||
capped_recall[f"R_cap@{k}"] = round(capped_recall[f"R_cap@{k}"]/len(qrels), 5)
|
||||
logging.info("R_cap@{}: {:.4f}".format(k, capped_recall[f"R_cap@{k}"]))
|
||||
|
||||
return capped_recall
|
||||
|
||||
|
||||
# Modified from https://github.com/embeddings-benchmark/mteb/blob/18f730696451a5aaa026494cecf288fd5cde9fd0/mteb/evaluation/evaluators/RetrievalEvaluator.py#L501
|
||||
def evaluate_metrics(
|
||||
qrels: Dict[str, Dict[str, int]],
|
||||
results: Dict[str, Dict[str, float]],
|
||||
k_values: List[int],
|
||||
) -> Tuple[
|
||||
Dict[str, float],
|
||||
Dict[str, float],
|
||||
Dict[str, float],
|
||||
Dict[str, float],
|
||||
]:
|
||||
"""Evaluate the main metrics.
|
||||
|
||||
Args:
|
||||
qrels (Dict[str, Dict[str, int]]): Ground truth relevance.
|
||||
results (Dict[str, Dict[str, float]]): Search results to evaluate.
|
||||
k_values (List[int]): Cutoffs.
|
||||
|
||||
Returns:
|
||||
Tuple[ Dict[str, float], Dict[str, float], Dict[str, float], Dict[str, float], ]: Results of different metrics at
|
||||
different provided k values.
|
||||
"""
|
||||
all_ndcgs, all_aps, all_recalls, all_precisions = defaultdict(list), defaultdict(list), defaultdict(list), defaultdict(list)
|
||||
|
||||
map_string = "map_cut." + ",".join([str(k) for k in k_values])
|
||||
ndcg_string = "ndcg_cut." + ",".join([str(k) for k in k_values])
|
||||
recall_string = "recall." + ",".join([str(k) for k in k_values])
|
||||
precision_string = "P." + ",".join([str(k) for k in k_values])
|
||||
evaluator = pytrec_eval.RelevanceEvaluator(
|
||||
qrels, {map_string, ndcg_string, recall_string, precision_string}
|
||||
)
|
||||
scores = evaluator.evaluate(results)
|
||||
|
||||
for query_id in scores.keys():
|
||||
for k in k_values:
|
||||
all_ndcgs[f"NDCG@{k}"].append(scores[query_id]["ndcg_cut_" + str(k)])
|
||||
all_aps[f"MAP@{k}"].append(scores[query_id]["map_cut_" + str(k)])
|
||||
all_recalls[f"Recall@{k}"].append(scores[query_id]["recall_" + str(k)])
|
||||
all_precisions[f"P@{k}"].append(scores[query_id]["P_" + str(k)])
|
||||
|
||||
ndcg, _map, recall, precision = (
|
||||
all_ndcgs.copy(),
|
||||
all_aps.copy(),
|
||||
all_recalls.copy(),
|
||||
all_precisions.copy(),
|
||||
)
|
||||
|
||||
for k in k_values:
|
||||
ndcg[f"NDCG@{k}"] = round(sum(ndcg[f"NDCG@{k}"]) / len(scores), 5)
|
||||
_map[f"MAP@{k}"] = round(sum(_map[f"MAP@{k}"]) / len(scores), 5)
|
||||
recall[f"Recall@{k}"] = round(sum(recall[f"Recall@{k}"]) / len(scores), 5)
|
||||
precision[f"P@{k}"] = round(sum(precision[f"P@{k}"]) / len(scores), 5)
|
||||
|
||||
return ndcg, _map, recall, precision
|
||||
|
||||
|
||||
def index(
|
||||
index_factory: str = "Flat",
|
||||
corpus_embeddings: Optional[np.ndarray] = None,
|
||||
load_path: Optional[str] = None,
|
||||
device: Optional[str] = None
|
||||
):
|
||||
"""Create and add embeddings into a Faiss index.
|
||||
|
||||
Args:
|
||||
index_factory (str, optional): Type of Faiss index to create. Defaults to "Flat".
|
||||
corpus_embeddings (Optional[np.ndarray], optional): The embedding vectors of the corpus. Defaults to None.
|
||||
load_path (Optional[str], optional): Path to load embeddings from. Defaults to None.
|
||||
device (Optional[str], optional): Device to hold Faiss index. Defaults to None.
|
||||
|
||||
Returns:
|
||||
faiss.Index: The Faiss index that contains all the corpus embeddings.
|
||||
"""
|
||||
if corpus_embeddings is None:
|
||||
corpus_embeddings = np.load(load_path)
|
||||
|
||||
logger.info(f"Shape of embeddings: {corpus_embeddings.shape}")
|
||||
# create faiss index
|
||||
logger.info(f'Indexing {corpus_embeddings.shape[0]} documents...')
|
||||
faiss_index = faiss.index_factory(corpus_embeddings.shape[-1], index_factory, faiss.METRIC_INNER_PRODUCT)
|
||||
|
||||
if device is None and torch.cuda.is_available():
|
||||
try:
|
||||
co = faiss.GpuMultipleClonerOptions()
|
||||
co.shard = True
|
||||
co.useFloat16 = True
|
||||
faiss_index = faiss.index_cpu_to_all_gpus(faiss_index, co)
|
||||
except:
|
||||
print('faiss do not support GPU, please uninstall faiss-cpu, faiss-gpu and install faiss-gpu again.')
|
||||
|
||||
logger.info('Adding embeddings ...')
|
||||
corpus_embeddings = corpus_embeddings.astype(np.float32)
|
||||
faiss_index.train(corpus_embeddings)
|
||||
faiss_index.add(corpus_embeddings)
|
||||
logger.info('Embeddings add over...')
|
||||
return faiss_index
|
||||
|
||||
|
||||
def search(
|
||||
faiss_index: faiss.Index,
|
||||
k: int = 100,
|
||||
query_embeddings: Optional[np.ndarray] = None,
|
||||
load_path: Optional[str] = None
|
||||
):
|
||||
"""
|
||||
1. Encode queries into dense embeddings;
|
||||
2. Search through faiss index
|
||||
|
||||
Args:
|
||||
faiss_index (faiss.Index): The Faiss index that contains all the corpus embeddings.
|
||||
k (int, optional): Top k numbers of closest neighbours. Defaults to :data:`100`.
|
||||
query_embeddings (Optional[np.ndarray], optional): The embedding vectors of queries. Defaults to :data:`None`.
|
||||
load_path (Optional[str], optional): Path to load embeddings from. Defaults to :data:`None`.
|
||||
|
||||
Returns:
|
||||
Tuple[np.ndarray, np.ndarray]: The scores of search results and their corresponding indices.
|
||||
"""
|
||||
if query_embeddings is None:
|
||||
query_embeddings = np.load(load_path)
|
||||
|
||||
query_size = len(query_embeddings)
|
||||
|
||||
all_scores = []
|
||||
all_indices = []
|
||||
|
||||
for i in tqdm(range(0, query_size, 32), desc="Searching"):
|
||||
j = min(i + 32, query_size)
|
||||
query_embedding = query_embeddings[i: j]
|
||||
score, indice = faiss_index.search(query_embedding.astype(np.float32), k=k)
|
||||
all_scores.append(score)
|
||||
all_indices.append(indice)
|
||||
|
||||
all_scores = np.concatenate(all_scores, axis=0)
|
||||
all_indices = np.concatenate(all_indices, axis=0)
|
||||
return all_scores, all_indices
|
||||
@@ -0,0 +1,143 @@
|
||||
import os
|
||||
from typing import Optional, List
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from transformers import TrainingArguments
|
||||
|
||||
|
||||
@dataclass
|
||||
class AbsEmbedderModelArguments:
|
||||
"""
|
||||
Abstract class for model arguments.
|
||||
"""
|
||||
|
||||
model_name_or_path: str = field(
|
||||
metadata={"help": "The model checkpoint for initialization."}
|
||||
)
|
||||
config_name: str = field(
|
||||
default=None,
|
||||
metadata={"help": "Pretrained config name or path if not the same as model_name."}
|
||||
)
|
||||
tokenizer_name: str = field(
|
||||
default=None,
|
||||
metadata={"help": "Pretrained tokenizer name or path if not the same as model_name."}
|
||||
)
|
||||
cache_dir: str = field(
|
||||
default=None,
|
||||
metadata={"help": "Where do you want to store the pre-trained models downloaded from s3."}
|
||||
)
|
||||
trust_remote_code: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "Trust remote code"}
|
||||
)
|
||||
use_fast_tokenizer: bool = field(
|
||||
default=True,
|
||||
metadata={"help": "Whether to use fast tokenizer or not."}
|
||||
)
|
||||
token: str = field(
|
||||
default_factory=lambda: os.getenv('HF_TOKEN', None),
|
||||
metadata={"help": "The token to use when accessing the model."}
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AbsEmbedderDataArguments:
|
||||
"""
|
||||
Abstract class for data arguments.
|
||||
"""
|
||||
train_data: str = field(
|
||||
default=None, metadata={
|
||||
"help": "One or more paths to training data. `query: str`, `pos: List[str]`, `neg: List[str]` are required in the training data.",
|
||||
"nargs": "+"
|
||||
}
|
||||
)
|
||||
cache_path: Optional[str] = field(
|
||||
default=None, metadata={"help": "Where do you want to store the cached data"}
|
||||
)
|
||||
train_group_size: int = field(default=8)
|
||||
|
||||
query_max_len: int = field(
|
||||
default=32,
|
||||
metadata={
|
||||
"help": "The maximum total input sequence length after tokenization for passage. Sequences longer than this will be truncated."
|
||||
},
|
||||
)
|
||||
|
||||
passage_max_len: int = field(
|
||||
default=128,
|
||||
metadata={
|
||||
"help": "The maximum total input sequence length after tokenization for passage. Sequences longer than this will be truncated."
|
||||
},
|
||||
)
|
||||
|
||||
pad_to_multiple_of: Optional[int] = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": "If set will pad the sequence to be a multiple of the provided value."
|
||||
},
|
||||
)
|
||||
|
||||
max_example_num_per_dataset: int = field(
|
||||
default=100000000, metadata={"help": "the max number of examples for each dataset"}
|
||||
)
|
||||
|
||||
query_instruction_for_retrieval: str= field(
|
||||
default=None, metadata={"help": "instruction for query"}
|
||||
)
|
||||
query_instruction_format: str = field(
|
||||
default="{}{}", metadata={"help": "format for query instruction"}
|
||||
)
|
||||
|
||||
knowledge_distillation: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "Use knowledge distillation when `pos_scores: List[float]` and `neg_scores: List[float]` are in features of training data"}
|
||||
)
|
||||
|
||||
passage_instruction_for_retrieval: Optional[str] = field(
|
||||
default=None, metadata={"help": "instruction for passage"}
|
||||
)
|
||||
passage_instruction_format: Optional[str] = field(
|
||||
default="{}{}", metadata={"help": "format for passage instruction"}
|
||||
)
|
||||
|
||||
shuffle_ratio: float = field(
|
||||
default=0.0, metadata={"help": "The ratio of shuffling the text"}
|
||||
)
|
||||
|
||||
# Parameters for SameDatasetDataArguments
|
||||
same_dataset_within_batch: bool = field(
|
||||
default=False, metadata={"help": "All samples in the same batch comes from the same dataset."}
|
||||
)
|
||||
small_threshold: int = field(
|
||||
default=0,
|
||||
metadata={"help": "The threshold of small dataset. All small dataset in the same directory will be merged into one dataset."}
|
||||
)
|
||||
drop_threshold: int = field(
|
||||
default=0,
|
||||
metadata={"help": "The threshold for dropping merged small dataset. If the number of examples in the merged small dataset is less than this threshold, it will be dropped."}
|
||||
)
|
||||
|
||||
def __post_init__(self):
|
||||
# replace "\\n" with "\n"
|
||||
if "\\n" in self.query_instruction_format:
|
||||
self.query_instruction_format = self.query_instruction_format.replace("\\n", "\n")
|
||||
if "\\n" in self.passage_instruction_format:
|
||||
self.passage_instruction_format = self.passage_instruction_format.replace("\\n", "\n")
|
||||
|
||||
# check the existence of train data
|
||||
for train_dir in self.train_data:
|
||||
if not os.path.exists(train_dir):
|
||||
raise FileNotFoundError(f"cannot find file: {train_dir}, please set a true path")
|
||||
|
||||
|
||||
@dataclass
|
||||
class AbsEmbedderTrainingArguments(TrainingArguments):
|
||||
negatives_cross_device: bool = field(default=False, metadata={"help": "share negatives across devices"})
|
||||
temperature: Optional[float] = field(default=0.02, metadata={"help": "temperature used for similarity score"})
|
||||
fix_position_embedding: bool = field(default=False, metadata={"help": "Freeze the parameters of position embeddings"})
|
||||
sentence_pooling_method: str = field(default='cls', metadata={"help": "the pooling method. Available options: cls, mean, last_token. Default: cls", "choices": ['cls', 'mean', 'last_token']})
|
||||
normalize_embeddings: bool = field(default=True, metadata={"help": "whether to normalize the embeddings"})
|
||||
sub_batch_size: Optional[int] = field(default=None, metadata={"help": "sub batch size for training"})
|
||||
kd_loss_type: str = field(default='kl_div', metadata={"help": "the loss type for knowledge distillation. Available options: kl_div, m3_kd_loss. Default: kl_div.", "choices": ['kl_div', 'm3_kd_loss']})
|
||||
use_mrl: bool = field(default=False, metadata={"help": "whether to use MRL for training"})
|
||||
mrl_dims: List[int] = field(default_factory=lambda: [], metadata={"help": "the dimensions of MRL layers"})
|
||||
@@ -0,0 +1,624 @@
|
||||
import os
|
||||
import math
|
||||
import random
|
||||
import logging
|
||||
import datasets
|
||||
import numpy as np
|
||||
import torch.distributed as dist
|
||||
from dataclasses import dataclass
|
||||
from torch.utils.data import Dataset
|
||||
from transformers import (
|
||||
PreTrainedTokenizer,
|
||||
DataCollatorWithPadding,
|
||||
TrainerCallback,
|
||||
TrainerState,
|
||||
TrainerControl
|
||||
)
|
||||
|
||||
from .AbsArguments import AbsEmbedderDataArguments, AbsEmbedderTrainingArguments
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsEmbedderTrainDataset(Dataset):
|
||||
"""Abstract class for training dataset.
|
||||
|
||||
Args:
|
||||
args (AbsEmbedderDataArguments): Data arguments.
|
||||
tokenizer (PreTrainedTokenizer): Tokenizer to use.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
args: AbsEmbedderDataArguments,
|
||||
tokenizer: PreTrainedTokenizer
|
||||
):
|
||||
self.args = args
|
||||
self.tokenizer = tokenizer
|
||||
self.shuffle_ratio = args.shuffle_ratio
|
||||
|
||||
train_datasets = []
|
||||
for data_dir in args.train_data:
|
||||
if not os.path.isdir(data_dir):
|
||||
if not (data_dir.endswith('.json') or data_dir.endswith('.jsonl')): continue
|
||||
temp_dataset = self._load_dataset(data_dir)
|
||||
if len(temp_dataset) == 0: continue
|
||||
train_datasets.append(temp_dataset)
|
||||
else:
|
||||
for file in os.listdir(data_dir):
|
||||
if not (file.endswith('.json') or file.endswith('.jsonl')): continue
|
||||
temp_dataset = self._load_dataset(os.path.join(data_dir, file))
|
||||
if len(temp_dataset) == 0: continue
|
||||
train_datasets.append(temp_dataset)
|
||||
self.dataset = datasets.concatenate_datasets(train_datasets)
|
||||
|
||||
def _load_dataset(self, file_path: str):
|
||||
"""Load dataset from path.
|
||||
|
||||
Args:
|
||||
file_path (str): Path to load the datasets from.
|
||||
|
||||
Raises:
|
||||
ValueError: `pos_scores` and `neg_scores` not found in the features of training data
|
||||
|
||||
Returns:
|
||||
datasets.Dataset: Loaded HF dataset.
|
||||
"""
|
||||
safe_rank = dist.get_rank() if dist.is_initialized() else 0
|
||||
if safe_rank == 0:
|
||||
logger.info(f'loading data from {file_path} ...')
|
||||
|
||||
temp_dataset = datasets.load_dataset('json', data_files=file_path, split='train', cache_dir=self.args.cache_path)
|
||||
if len(temp_dataset) > self.args.max_example_num_per_dataset:
|
||||
temp_dataset = temp_dataset.select(random.sample(list(range(len(temp_dataset))), self.args.max_example_num_per_dataset))
|
||||
if not self.args.knowledge_distillation:
|
||||
if 'pos_scores' in temp_dataset.column_names:
|
||||
temp_dataset = temp_dataset.remove_columns(['pos_scores'])
|
||||
if 'neg_scores' in temp_dataset.column_names:
|
||||
temp_dataset = temp_dataset.remove_columns(['neg_scores'])
|
||||
else:
|
||||
if 'pos_scores' not in temp_dataset.column_names or 'neg_scores' not in temp_dataset.column_names:
|
||||
raise ValueError(f"`pos_scores` and `neg_scores` not found in the features of training data in {file_path}, which is necessary when using knowledge distillation.")
|
||||
return temp_dataset
|
||||
|
||||
def _shuffle_text(self, text):
|
||||
"""shuffle the input text.
|
||||
|
||||
Args:
|
||||
text (str): Input text.
|
||||
|
||||
Returns:
|
||||
str: Shuffled text.
|
||||
"""
|
||||
if self.shuffle_ratio > 0 and len(text) > 100 and random.random() < self.shuffle_ratio:
|
||||
split_text = []
|
||||
chunk_size = len(text)//3 + 1
|
||||
for i in range(0, len(text), chunk_size):
|
||||
split_text.append(text[i:i+chunk_size])
|
||||
random.shuffle(split_text)
|
||||
return " ".join(split_text)
|
||||
else:
|
||||
return text
|
||||
|
||||
def __len__(self):
|
||||
return len(self.dataset)
|
||||
|
||||
def __getitem__(self, item):
|
||||
data = self.dataset[item]
|
||||
train_group_size = self.args.train_group_size
|
||||
|
||||
query = data['query']
|
||||
if self.args.query_instruction_for_retrieval is not None:
|
||||
query = self.args.query_instruction_format.format(
|
||||
data['prompt'] if 'prompt' in data else self.args.query_instruction_for_retrieval,
|
||||
query
|
||||
)
|
||||
|
||||
passages = []
|
||||
teacher_scores = []
|
||||
|
||||
assert isinstance(data['pos'], list) and isinstance(data['neg'], list)
|
||||
|
||||
pos_idx = random.choice(list(range(len(data['pos']))))
|
||||
passages.append(self._shuffle_text(data['pos'][pos_idx]))
|
||||
|
||||
neg_all_idx = list(range(len(data['neg'])))
|
||||
if len(data['neg']) < train_group_size - 1:
|
||||
num = math.ceil((train_group_size - 1) / len(data['neg']))
|
||||
neg_idxs = random.sample(neg_all_idx * num, train_group_size - 1)
|
||||
else:
|
||||
neg_idxs = random.sample(neg_all_idx, self.args.train_group_size - 1)
|
||||
for neg_idx in neg_idxs:
|
||||
passages.append(data['neg'][neg_idx])
|
||||
|
||||
if self.args.knowledge_distillation:
|
||||
assert isinstance(data['pos_scores'], list) and isinstance(data['neg_scores'], list)
|
||||
teacher_scores.append(data['pos_scores'][pos_idx])
|
||||
for neg_idx in neg_idxs:
|
||||
teacher_scores.append(data['neg_scores'][neg_idx])
|
||||
if not all(isinstance(score, (int, float)) for score in teacher_scores):
|
||||
raise ValueError(f"pos_score or neg_score must be digit")
|
||||
else:
|
||||
teacher_scores = None
|
||||
|
||||
if self.args.passage_instruction_for_retrieval is not None:
|
||||
passages = [
|
||||
self.args.passage_instruction_format.format(
|
||||
self.args.passage_instruction_for_retrieval, p
|
||||
)
|
||||
for p in passages
|
||||
]
|
||||
|
||||
return query, passages, teacher_scores
|
||||
|
||||
@dataclass
|
||||
class AbsEmbedderCollator(DataCollatorWithPadding):
|
||||
"""
|
||||
The abstract embedder collator.
|
||||
"""
|
||||
query_max_len: int = 32
|
||||
passage_max_len: int = 128
|
||||
sub_batch_size: int = -1
|
||||
|
||||
def __call__(self, features):
|
||||
queries = [f[0] for f in features]
|
||||
passages = [f[1] for f in features]
|
||||
teacher_scores = [f[2] for f in features]
|
||||
if teacher_scores[0] is None:
|
||||
teacher_scores = None
|
||||
elif isinstance(teacher_scores[0], list):
|
||||
teacher_scores = sum(teacher_scores, [])
|
||||
|
||||
if isinstance(queries[0], list):
|
||||
queries = sum(queries, [])
|
||||
if isinstance(passages[0], list):
|
||||
passages = sum(passages, [])
|
||||
|
||||
queries_inputs = self.tokenizer(
|
||||
queries,
|
||||
truncation=True,
|
||||
max_length=self.query_max_len,
|
||||
return_tensors=None
|
||||
)
|
||||
passages_inputs = self.tokenizer(
|
||||
passages,
|
||||
truncation=True,
|
||||
max_length=self.passage_max_len,
|
||||
return_tensors=None
|
||||
)
|
||||
|
||||
if self.sub_batch_size is None or self.sub_batch_size <= 0:
|
||||
q_collated = self.tokenizer.pad(
|
||||
queries_inputs,
|
||||
padding=self.padding,
|
||||
max_length=self.query_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=self.return_tensors
|
||||
)
|
||||
d_collated = self.tokenizer.pad(
|
||||
passages_inputs,
|
||||
padding=self.padding,
|
||||
max_length=self.passage_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=self.return_tensors
|
||||
)
|
||||
else:
|
||||
batch_size = self.sub_batch_size
|
||||
|
||||
q_collated = []
|
||||
for i in range(0, len(queries_inputs['attention_mask']), batch_size):
|
||||
start = i
|
||||
end = min(len(queries_inputs['attention_mask']), i + batch_size)
|
||||
sub_features = {}
|
||||
for k, v in queries_inputs.items():
|
||||
sub_features[k] = v[start:end]
|
||||
q_collated.append(self.tokenizer.pad(
|
||||
sub_features,
|
||||
padding=self.padding,
|
||||
max_length=self.query_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=self.return_tensors
|
||||
))
|
||||
|
||||
d_collated = []
|
||||
for i in range(0, len(passages_inputs['attention_mask']), batch_size):
|
||||
start = i
|
||||
end = min(len(passages_inputs['attention_mask']), i + batch_size)
|
||||
sub_features = {}
|
||||
|
||||
for k, v in passages_inputs.items():
|
||||
sub_features[k] = v[start:end]
|
||||
d_collated.append(self.tokenizer.pad(
|
||||
sub_features,
|
||||
padding=self.padding,
|
||||
max_length=self.passage_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=self.return_tensors
|
||||
))
|
||||
return {
|
||||
"queries": q_collated,
|
||||
"passages": d_collated,
|
||||
"teacher_scores": teacher_scores,
|
||||
"no_in_batch_neg_flag": False
|
||||
}
|
||||
|
||||
|
||||
class AbsEmbedderSameDatasetTrainDataset(AbsEmbedderTrainDataset):
|
||||
"""Abstract class for training dataset that samples batches from same dataset.
|
||||
|
||||
Args:
|
||||
args (AbsEmbedderDataArguments): Data arguments.
|
||||
default_batch_size (int): The default batch size for training.
|
||||
seed (int): Random seed.
|
||||
tokenizer (PreTrainedTokenizer): Tokenizer to use.
|
||||
process_index (int, optional): Current process index. Defaults to 0.
|
||||
num_processes (int, optional): Total number of processes. Defaults to 1.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
args: AbsEmbedderDataArguments,
|
||||
default_batch_size: int,
|
||||
seed: int,
|
||||
tokenizer: PreTrainedTokenizer,
|
||||
process_index: int=0,
|
||||
num_processes: int=1
|
||||
):
|
||||
self.args = args
|
||||
self.shuffle_ratio = args.shuffle_ratio
|
||||
self.defaut_batch_size = default_batch_size
|
||||
self.deterministic_generator = np.random.default_rng(seed)
|
||||
self.tokenizer = tokenizer
|
||||
self.process_index = process_index
|
||||
self.num_processes = num_processes
|
||||
|
||||
self.step = 0
|
||||
|
||||
train_datasets = []
|
||||
each_data_idxs = []
|
||||
batch_size_idxs = []
|
||||
no_in_batch_neg_flags = []
|
||||
cur_all_num = 0
|
||||
|
||||
small_threshold = args.small_threshold
|
||||
drop_threshold = args.drop_threshold
|
||||
|
||||
for data_dir in args.train_data:
|
||||
if not os.path.isdir(data_dir):
|
||||
# Add `no_in_batch_neg` **suffix** to `data_dir` to indicate that this dataset does not use in-batch negatives
|
||||
no_in_batch_neg_flag = data_dir.split('.')[-2].endswith('no_in_batch_neg')
|
||||
if not (data_dir.endswith('.json') or data_dir.endswith('.jsonl')): continue
|
||||
temp_dataset = self._load_dataset(data_dir)
|
||||
|
||||
if len(temp_dataset) == 0 or len(temp_dataset) < small_threshold: continue
|
||||
else:
|
||||
train_datasets.append(temp_dataset)
|
||||
each_data_idxs.append(np.arange(len(temp_dataset)) + cur_all_num)
|
||||
cur_all_num += len(temp_dataset)
|
||||
batch_size_idxs.append(self._get_file_batch_size(temp_dataset, default_batch_size))
|
||||
no_in_batch_neg_flags.append(no_in_batch_neg_flag)
|
||||
|
||||
else:
|
||||
small_datasets = []
|
||||
small_batch_size = math.inf
|
||||
|
||||
# Add `no_in_batch_neg` **suffix** to `data_dir` to indicate that this dataset does not use in-batch negatives
|
||||
no_in_batch_neg_flag = data_dir.endswith('no_in_batch_neg')
|
||||
for file in os.listdir(data_dir):
|
||||
if not (file.endswith('.json') or file.endswith('.jsonl')): continue
|
||||
temp_dataset = self._load_dataset(os.path.join(data_dir, file))
|
||||
|
||||
if len(temp_dataset) == 0: continue
|
||||
elif len(temp_dataset) < small_threshold:
|
||||
small_datasets.append(temp_dataset)
|
||||
small_batch_size = min(small_batch_size, self._get_file_batch_size(temp_dataset, default_batch_size))
|
||||
else:
|
||||
train_datasets.append(temp_dataset)
|
||||
each_data_idxs.append(np.arange(len(temp_dataset)) + cur_all_num)
|
||||
cur_all_num += len(temp_dataset)
|
||||
batch_size_idxs.append(self._get_file_batch_size(temp_dataset, default_batch_size))
|
||||
no_in_batch_neg_flags.append(no_in_batch_neg_flag)
|
||||
|
||||
if len(small_datasets) > 0:
|
||||
small_dataset = datasets.concatenate_datasets(small_datasets)
|
||||
if len(small_dataset) >= drop_threshold:
|
||||
train_datasets.append(small_dataset)
|
||||
each_data_idxs.append(np.arange(len(small_dataset)) + cur_all_num)
|
||||
cur_all_num += len(small_dataset)
|
||||
batch_size_idxs.append(small_batch_size)
|
||||
no_in_batch_neg_flags.append(no_in_batch_neg_flag)
|
||||
|
||||
self.dataset = datasets.concatenate_datasets(train_datasets)
|
||||
self.each_data_idxs = each_data_idxs
|
||||
self.datasets_inxs = np.arange(len(each_data_idxs))
|
||||
self.batch_size_idxs = batch_size_idxs
|
||||
self.no_in_batch_neg_flags = no_in_batch_neg_flags
|
||||
|
||||
self.refresh_epoch()
|
||||
|
||||
def _load_dataset(self, file_path: str):
|
||||
"""Load datset from given path.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to load or download from HF hub.
|
||||
|
||||
Returns:
|
||||
datasets.Dataset: The loaded dataset.
|
||||
"""
|
||||
safe_rank = dist.get_rank() if dist.is_initialized() else 0
|
||||
if safe_rank == 0:
|
||||
logger.info(f'loading data from {file_path} ...')
|
||||
|
||||
temp_dataset = datasets.load_dataset('json', data_files=file_path, split='train', cache_dir=self.args.cache_path)
|
||||
if len(temp_dataset) > self.args.max_example_num_per_dataset:
|
||||
temp_dataset = temp_dataset.select(random.sample(list(range(len(temp_dataset))), self.args.max_example_num_per_dataset))
|
||||
if not self.args.knowledge_distillation:
|
||||
if 'pos_scores' in temp_dataset.column_names:
|
||||
temp_dataset = temp_dataset.remove_columns(['pos_scores'])
|
||||
if 'neg_scores' in temp_dataset.column_names:
|
||||
temp_dataset = temp_dataset.remove_columns(['neg_scores'])
|
||||
return temp_dataset
|
||||
|
||||
@staticmethod
|
||||
def _get_file_batch_size(temp_dataset: datasets.Dataset, default_batch_size: int):
|
||||
"""Get the appropriate batch size for the dataset.
|
||||
|
||||
Args:
|
||||
temp_dataset (datasets.Dataset): Loaded :data:`datasets.Dataset` object.
|
||||
default_batch_size (int): The default batch size to use if not specified in the dataset.
|
||||
|
||||
Returns:
|
||||
int: The final batch size to use.
|
||||
"""
|
||||
if 'batch_size' in temp_dataset.column_names:
|
||||
return temp_dataset['batch_size'][0]
|
||||
if 'type' in temp_dataset.column_names:
|
||||
data_type = temp_dataset['type'][0]
|
||||
if 'symmetric' in data_type:
|
||||
return default_batch_size // 2 # make the symmetric data have smaller batch size
|
||||
return default_batch_size
|
||||
|
||||
def refresh_epoch(self):
|
||||
"""
|
||||
Refresh data for epoch.
|
||||
"""
|
||||
logger.info(f'-- Rank {self.process_index}: refresh data --')
|
||||
self.deterministic_generator.shuffle(self.datasets_inxs)
|
||||
|
||||
batch_datas = []
|
||||
for dataset_inx in self.datasets_inxs:
|
||||
self.deterministic_generator.shuffle(self.each_data_idxs[dataset_inx])
|
||||
cur_batch_size = self.batch_size_idxs[dataset_inx]*self.num_processes
|
||||
no_in_batch_neg_flag = self.no_in_batch_neg_flags[dataset_inx]
|
||||
for start_index in range(0, len(self.each_data_idxs[dataset_inx]), cur_batch_size):
|
||||
# judge the last batch's length
|
||||
if len(self.each_data_idxs[dataset_inx]) - start_index < cur_batch_size:
|
||||
break
|
||||
batch_datas.append((
|
||||
self.each_data_idxs[dataset_inx][start_index:start_index+cur_batch_size],
|
||||
no_in_batch_neg_flag
|
||||
))
|
||||
self.deterministic_generator.shuffle(batch_datas)
|
||||
self.batch_datas = batch_datas
|
||||
self.step = 0
|
||||
|
||||
def __len__(self):
|
||||
return len(self.batch_datas) * self.num_processes
|
||||
|
||||
def __getitem__(self, _):
|
||||
batch_indices, no_in_batch_neg_flag = self.batch_datas[self.step] # extend here
|
||||
cur_batch_size = int(len(batch_indices) / self.num_processes)
|
||||
batch_indices = batch_indices[self.process_index * cur_batch_size: (self.process_index + 1) * cur_batch_size]
|
||||
batch_data = self.dataset[batch_indices]
|
||||
self.step += 1
|
||||
queries, passages, teacher_scores = self._create_batch_data(batch_raw_data=batch_data)
|
||||
return queries, passages, teacher_scores, no_in_batch_neg_flag
|
||||
|
||||
def _get_train_group_size(self, batch_raw_data):
|
||||
"""Get the training group size and data type.
|
||||
|
||||
Args:
|
||||
batch_raw_data (datasets.Dataset): One batch of raw data.
|
||||
|
||||
Returns:
|
||||
int: The training group size.
|
||||
str: The type of data for the task.
|
||||
"""
|
||||
if 'type' in batch_raw_data:
|
||||
data_type = batch_raw_data['type'][0]
|
||||
if data_type in ['only_1neg']:
|
||||
return 2, data_type
|
||||
elif data_type in ['symmetric_class']:
|
||||
return min(len(batch_raw_data['neg'][0]) + 1, self.args.train_group_size), data_type
|
||||
else:
|
||||
return self.args.train_group_size, data_type
|
||||
elif 'train_group_size' in batch_raw_data:
|
||||
train_group_size = batch_raw_data['train_group_size'][0]
|
||||
if isinstance(train_group_size, int) and train_group_size > 0:
|
||||
return train_group_size, None
|
||||
else:
|
||||
return self.args.train_group_size, None
|
||||
return self.args.train_group_size, None
|
||||
|
||||
def _create_batch_data(self, batch_raw_data):
|
||||
"""Create a comple batch of data with queries, documents and teacher scores.
|
||||
|
||||
Args:
|
||||
batch_raw_data (datasets.Dataset): One batch of raw data.
|
||||
|
||||
Returns:
|
||||
List[str]: Queries with instruction format.
|
||||
List[str]: Documents with instruction format.
|
||||
List[float]: Teacher scores for model distillation.
|
||||
"""
|
||||
queries, passages, teacher_scores = [], [], []
|
||||
|
||||
train_group_size, data_type = self._get_train_group_size(batch_raw_data)
|
||||
|
||||
for i in range(len(batch_raw_data['query'])):
|
||||
if data_type is not None:
|
||||
assert batch_raw_data['type'][i] == data_type, f"Data type is not consistent in the same batch"
|
||||
|
||||
queries.append(
|
||||
self.args.query_instruction_format.format(
|
||||
batch_raw_data['prompt'][i] if 'prompt' in batch_raw_data else self.args.query_instruction_for_retrieval,
|
||||
batch_raw_data['query'][i]
|
||||
)
|
||||
)
|
||||
tmp_passages = []
|
||||
pos_idx = random.choice(list(range(len(batch_raw_data['pos'][i]))))
|
||||
pos = self._shuffle_text(batch_raw_data['pos'][i][pos_idx])
|
||||
tmp_passages.append(pos)
|
||||
|
||||
neg_all_idx = list(range(len(batch_raw_data['neg'][i])))
|
||||
if len(batch_raw_data['neg'][i]) < train_group_size - 1:
|
||||
num = math.ceil((train_group_size - 1) / len(batch_raw_data['neg'][i]))
|
||||
neg_idxs = random.sample(neg_all_idx * num, train_group_size - 1)
|
||||
else:
|
||||
neg_idxs = random.sample(neg_all_idx, train_group_size - 1)
|
||||
for neg_idx in neg_idxs:
|
||||
tmp_passages.append(batch_raw_data['neg'][i][neg_idx])
|
||||
|
||||
if self.args.knowledge_distillation:
|
||||
if 'pos_scores' in batch_raw_data and batch_raw_data['pos_scores'][i] is not None:
|
||||
teacher_scores.append(batch_raw_data['pos_scores'][i][pos_idx])
|
||||
for neg_idx in neg_idxs:
|
||||
if 'neg_scores' in batch_raw_data and batch_raw_data['neg_scores'][i] is not None:
|
||||
teacher_scores.append(batch_raw_data['neg_scores'][i][neg_idx])
|
||||
else:
|
||||
teacher_scores = None
|
||||
|
||||
if data_type is not None and data_type in ['symmetric_sts', 'symmetric_clustering']:
|
||||
tmp_passages = [
|
||||
self.args.query_instruction_format.format(
|
||||
batch_raw_data['prompt'][i] if 'prompt' in batch_raw_data else self.args.query_instruction_for_retrieval,
|
||||
p
|
||||
) for p in tmp_passages
|
||||
]
|
||||
else:
|
||||
if self.args.passage_instruction_for_retrieval is not None:
|
||||
tmp_passages = [
|
||||
self.args.passage_instruction_format.format(
|
||||
self.args.passage_instruction_for_retrieval, p
|
||||
) for p in tmp_passages
|
||||
]
|
||||
|
||||
passages.extend(tmp_passages)
|
||||
|
||||
if teacher_scores is not None:
|
||||
if len(teacher_scores) > 0 and len(passages) > 0:
|
||||
assert len(teacher_scores) == len(passages)
|
||||
|
||||
return queries, passages, teacher_scores
|
||||
|
||||
|
||||
@dataclass
|
||||
class AbsEmbedderSameDatasetCollator(DataCollatorWithPadding):
|
||||
"""
|
||||
EmbedCollator for SameDataset.
|
||||
Note that after using this collator, the training_args should be set as:
|
||||
|
||||
``training_args.per_device_train_batch_size = 1``
|
||||
|
||||
``training_args.dataloader_num_workers = 0 # avoid multi-processing``
|
||||
"""
|
||||
query_max_len: int = 32
|
||||
passage_max_len: int = 128
|
||||
sub_batch_size: int = -1
|
||||
|
||||
def __call__(self, features):
|
||||
queries = features[0][0]
|
||||
passages = features[0][1]
|
||||
teacher_scores = features[0][2]
|
||||
no_in_batch_neg_flag = features[0][3]
|
||||
|
||||
queries_inputs = self.tokenizer(
|
||||
queries,
|
||||
truncation=True,
|
||||
max_length=self.query_max_len,
|
||||
return_tensors=None
|
||||
)
|
||||
passages_inputs = self.tokenizer(
|
||||
passages,
|
||||
truncation=True,
|
||||
max_length=self.passage_max_len,
|
||||
return_tensors=None
|
||||
)
|
||||
|
||||
if self.sub_batch_size is None or self.sub_batch_size <= 0:
|
||||
q_collated = self.tokenizer.pad(
|
||||
queries_inputs,
|
||||
padding=self.padding,
|
||||
max_length=self.query_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=self.return_tensors,
|
||||
)
|
||||
|
||||
d_collated = self.tokenizer.pad(
|
||||
passages_inputs,
|
||||
padding=self.padding,
|
||||
max_length=self.passage_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=self.return_tensors,
|
||||
)
|
||||
else:
|
||||
batch_size = self.sub_batch_size
|
||||
|
||||
q_collated = []
|
||||
for i in range(0, len(queries_inputs['attention_mask']), batch_size):
|
||||
start = i
|
||||
end = min(len(queries_inputs['attention_mask']), i + batch_size)
|
||||
sub_features = {}
|
||||
for k, v in queries_inputs.items():
|
||||
sub_features[k] = v[start:end]
|
||||
q_collated.append(self.tokenizer.pad(
|
||||
sub_features,
|
||||
padding=self.padding,
|
||||
max_length=self.query_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=self.return_tensors,
|
||||
))
|
||||
|
||||
d_collated = []
|
||||
for i in range(0, len(passages_inputs['attention_mask']), batch_size):
|
||||
start = i
|
||||
end = min(len(passages_inputs['attention_mask']), i + batch_size)
|
||||
sub_features = {}
|
||||
|
||||
for k, v in passages_inputs.items():
|
||||
sub_features[k] = v[start:end]
|
||||
d_collated.append(self.tokenizer.pad(
|
||||
sub_features,
|
||||
padding=self.padding,
|
||||
max_length=self.passage_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=self.return_tensors,
|
||||
))
|
||||
|
||||
if isinstance(teacher_scores, list) and len(teacher_scores) == 0:
|
||||
teacher_scores = None
|
||||
|
||||
return {
|
||||
"queries": q_collated,
|
||||
"passages": d_collated,
|
||||
"teacher_scores": teacher_scores,
|
||||
"no_in_batch_neg_flag": no_in_batch_neg_flag
|
||||
}
|
||||
|
||||
|
||||
class EmbedderTrainerCallbackForDataRefresh(TrainerCallback):
|
||||
"""
|
||||
Callback class to inspect the state of the training loop and take decision.
|
||||
"""
|
||||
def __init__(self, train_dataset: AbsEmbedderSameDatasetTrainDataset):
|
||||
self.train_dataset = train_dataset
|
||||
|
||||
def on_epoch_end(
|
||||
self,
|
||||
args: AbsEmbedderTrainingArguments,
|
||||
state: TrainerState,
|
||||
control: TrainerControl,
|
||||
**kwargs
|
||||
):
|
||||
"""
|
||||
Event called at the end of an epoch.
|
||||
"""
|
||||
self.train_dataset.refresh_epoch()
|
||||
@@ -0,0 +1,364 @@
|
||||
import torch
|
||||
from torch import nn, Tensor
|
||||
import torch.nn.functional as F
|
||||
import torch.distributed as dist
|
||||
from transformers import PreTrainedTokenizer
|
||||
from transformers.file_utils import ModelOutput
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, Optional, List, Union
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class EmbedderOutput(ModelOutput):
|
||||
"""
|
||||
Output information returned by the model.
|
||||
"""
|
||||
q_reps: Optional[Tensor] = None
|
||||
p_reps: Optional[Tensor] = None
|
||||
loss: Optional[Tensor] = None
|
||||
scores: Optional[Tensor] = None
|
||||
|
||||
|
||||
class AbsEmbedderModel(ABC, nn.Module):
|
||||
"""Abstract class of embedding model for training.
|
||||
|
||||
Args:
|
||||
base_model: The base model to train on.
|
||||
tokenizer (PreTrainedTokenizer, optional): The tokenizer to use. Defaults to ``None``.
|
||||
negatives_cross_device (bool, optional): If True, will compute cross devices negative loss. Defaults to ``False``.
|
||||
temperature (float, optional): Temperature to control the scale of scores. Defaults to ``1.0``.
|
||||
sub_batch_size (int, optional): Sub-batch size during encoding. If negative, will not split to sub-batch.
|
||||
Defaults to ``-1``.
|
||||
kd_loss_type (str, optional): Type of knowledge distillation loss. Defaults to ``"kl_div"``.
|
||||
use_mrl (bool, optional): Whether to use MRL for training. Defaults to ``False``.
|
||||
mrl_dims (List[int], optional): The dimensions of MRL layers. Defaults to ``[]``.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
base_model,
|
||||
tokenizer: PreTrainedTokenizer = None,
|
||||
negatives_cross_device: bool = False,
|
||||
temperature: float = 1.0,
|
||||
sub_batch_size: int = -1,
|
||||
kd_loss_type: str = 'kl_div',
|
||||
use_mrl: bool = False,
|
||||
mrl_dims: List[int] = [],
|
||||
):
|
||||
nn.Module.__init__(self)
|
||||
self.model = base_model
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
self.temperature = temperature
|
||||
self.negatives_cross_device = negatives_cross_device
|
||||
if self.negatives_cross_device:
|
||||
if not dist.is_initialized():
|
||||
raise ValueError('Distributed training has not been initialized for representation all gather.')
|
||||
self.process_rank = dist.get_rank() if dist.is_initialized() else 0
|
||||
self.world_size = dist.get_world_size() if dist.is_initialized() else 1
|
||||
|
||||
self.sub_batch_size = sub_batch_size
|
||||
self.kd_loss_type = kd_loss_type
|
||||
|
||||
self.use_mrl = use_mrl
|
||||
self.mrl_dims = mrl_dims
|
||||
if self.use_mrl and len(self.mrl_dims) == 0:
|
||||
raise ValueError("mrl_dims should be provided when use_mrl is True")
|
||||
|
||||
@abstractmethod
|
||||
def encode(self, features):
|
||||
"""Abstract method encode and get the embedding.
|
||||
|
||||
Args:
|
||||
features (Union[list, dict]): Features feed to the model.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def compute_loss(self, scores, target):
|
||||
"""Abstract method compute the loss.
|
||||
|
||||
Args:
|
||||
scores (torch.Tensor): Computed score.
|
||||
target (torch.Tensor): The target value.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def compute_score(self, q_reps, p_reps):
|
||||
"""Abstract method to compute the score.
|
||||
|
||||
Args:
|
||||
q_reps (torch.Tensor): Queries representations.
|
||||
p_reps (torch.Tensor): Passages rerpresentations.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def save(self, output_dir: str):
|
||||
"""Abstract method to save the model.
|
||||
|
||||
Args:
|
||||
output_dir (str): Directory for saving the model.
|
||||
"""
|
||||
pass
|
||||
|
||||
def get_local_score(self, q_reps, p_reps, all_scores):
|
||||
"""Get the local score of queries and passages.
|
||||
|
||||
Args:
|
||||
q_reps (torch.Tensor): Queries representations.
|
||||
p_reps (torch.Tensor): Passages rerpresentations.
|
||||
all_scores (torch.Tensor): All the query-passage scores computed.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Local scores to compute loss.
|
||||
"""
|
||||
group_size = p_reps.size(0) // q_reps.size(0)
|
||||
indices = torch.arange(0, q_reps.size(0), device=q_reps.device) * group_size
|
||||
specific_scores = []
|
||||
for i in range(group_size):
|
||||
specific_scores.append(
|
||||
all_scores[torch.arange(q_reps.size(0), device=q_reps.device), indices + i]
|
||||
)
|
||||
return torch.stack(specific_scores, dim=1).view(q_reps.size(0), -1)
|
||||
|
||||
def compute_local_score(self, q_reps, p_reps, compute_score_func=None, **kwargs):
|
||||
"""Compute the local score of queries and passages.
|
||||
|
||||
Args:
|
||||
q_reps (torch.Tensor): Queries representations.
|
||||
p_reps (torch.Tensor): Passages rerpresentations.
|
||||
compute_score_func (function, optional): Function to compute score. Defaults to ``None``, which will use the
|
||||
:meth:`self.compute_score`.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Local scores to compute loss.
|
||||
"""
|
||||
if compute_score_func is None:
|
||||
all_scores = self.compute_score(q_reps, p_reps)
|
||||
else:
|
||||
all_scores = compute_score_func(q_reps, p_reps, **kwargs)
|
||||
loacl_scores = self.get_local_score(q_reps, p_reps, all_scores)
|
||||
return loacl_scores
|
||||
|
||||
def _compute_no_in_batch_neg_loss(self, q_reps, p_reps, teacher_targets=None, compute_score_func=None, **kwargs):
|
||||
"""
|
||||
Compute loss when using no in-batch negatives and no cross-device negatives
|
||||
"""
|
||||
group_size = p_reps.size(0) // q_reps.size(0)
|
||||
|
||||
local_scores = self.compute_local_score(q_reps, p_reps, compute_score_func, **kwargs) # (batch_size, group_size)
|
||||
|
||||
if teacher_targets is not None:
|
||||
# compute kd loss
|
||||
loss = self.distill_loss(self.kd_loss_type, teacher_targets, local_scores, group_size=group_size)
|
||||
|
||||
# add normal loss if needed
|
||||
if self.kd_loss_type == "kl_div":
|
||||
local_targets = torch.zeros(local_scores.size(0), device=local_scores.device, dtype=torch.long) # (batch_size)
|
||||
loss += self.compute_loss(local_scores, local_targets)
|
||||
else:
|
||||
local_targets = torch.zeros(local_scores.size(0), device=local_scores.device, dtype=torch.long) # (batch_size)
|
||||
loss = self.compute_loss(local_scores, local_targets)
|
||||
|
||||
return local_scores, loss
|
||||
|
||||
def _compute_in_batch_neg_loss(self, q_reps, p_reps, teacher_targets=None, compute_score_func=None, **kwargs):
|
||||
"""
|
||||
Compute loss when only using in-batch negatives
|
||||
"""
|
||||
group_size = p_reps.size(0) // q_reps.size(0)
|
||||
|
||||
if compute_score_func is None:
|
||||
scores = self.compute_score(q_reps, p_reps) # (batch_size, batch_size * group_size)
|
||||
else:
|
||||
scores = compute_score_func(q_reps, p_reps, **kwargs) # (batch_size, batch_size * group_size)
|
||||
|
||||
if teacher_targets is not None:
|
||||
# compute kd loss
|
||||
if self.kd_loss_type == "kl_div":
|
||||
student_scores = self.get_local_score(q_reps, p_reps, scores) # (batch_size, group_size)
|
||||
|
||||
loss = self.distill_loss(self.kd_loss_type, teacher_targets, student_scores, group_size)
|
||||
|
||||
idxs = torch.arange(q_reps.size(0), device=q_reps.device, dtype=torch.long)
|
||||
targets = idxs * (p_reps.size(0) // q_reps.size(0)) # (batch_size)
|
||||
loss += self.compute_loss(scores, targets)
|
||||
elif self.kd_loss_type == "m3_kd_loss":
|
||||
loss = self.distill_loss(self.kd_loss_type, teacher_targets, scores, group_size)
|
||||
else:
|
||||
raise ValueError(f"Invalid kd_loss_type: {self.kd_loss_type}")
|
||||
else:
|
||||
idxs = torch.arange(q_reps.size(0), device=q_reps.device, dtype=torch.long)
|
||||
targets = idxs * group_size # (batch_size)
|
||||
loss = self.compute_loss(scores, targets)
|
||||
|
||||
return scores, loss
|
||||
|
||||
def _compute_cross_device_neg_loss(self, q_reps, p_reps, teacher_targets=None, compute_score_func=None, **kwargs):
|
||||
"""
|
||||
Compute loss when using both in-batch negatives and cross-device negatives
|
||||
"""
|
||||
group_size = p_reps.size(0) // q_reps.size(0)
|
||||
|
||||
cross_q_reps = self._dist_gather_tensor(q_reps) # (world_size * batch_size, dim)
|
||||
cross_p_reps = self._dist_gather_tensor(p_reps) # (world_size * batch_size * group_size, dim)
|
||||
|
||||
if compute_score_func is None:
|
||||
cross_scores = self.compute_score(cross_q_reps, cross_p_reps) # (world_size * batch_size, world_size * batch_size * group_size)
|
||||
else:
|
||||
cross_scores = compute_score_func(cross_q_reps, cross_p_reps, **kwargs) # (world_size * batch_size, world_size * batch_size * group_size)
|
||||
|
||||
if teacher_targets is not None:
|
||||
# compute kd loss
|
||||
if self.kd_loss_type == "kl_div":
|
||||
student_scores = self.get_local_score(cross_q_reps, cross_p_reps, cross_scores) # (world_size * batch_size, group_size)
|
||||
student_scores = student_scores[
|
||||
q_reps.size(0)*self.process_rank : q_reps.size(0)*(self.process_rank+1)
|
||||
] # (batch_size, group_size)
|
||||
|
||||
loss = self.distill_loss(self.kd_loss_type, teacher_targets, student_scores, group_size)
|
||||
|
||||
cross_idxs = torch.arange(cross_q_reps.size(0), device=cross_q_reps.device, dtype=torch.long)
|
||||
cross_targets = cross_idxs * group_size # (world_size * batch_size)
|
||||
loss += self.compute_loss(cross_scores, cross_targets)
|
||||
elif self.kd_loss_type == "m3_kd_loss":
|
||||
cross_teacher_targets = self._dist_gather_tensor(teacher_targets) # (world_size * batch_size, group_size)
|
||||
|
||||
loss = self.distill_loss(self.kd_loss_type, cross_teacher_targets, cross_scores, group_size)
|
||||
else:
|
||||
raise ValueError(f"Invalid kd_loss_type: {self.kd_loss_type}")
|
||||
else:
|
||||
cross_idxs = torch.arange(cross_q_reps.size(0), device=cross_q_reps.device, dtype=torch.long)
|
||||
cross_targets = cross_idxs * group_size # (world_size * batch_size)
|
||||
loss = self.compute_loss(cross_scores, cross_targets)
|
||||
|
||||
return cross_scores, loss
|
||||
|
||||
def forward(
|
||||
self,
|
||||
queries: Union[Dict[str, Tensor], List[Dict[str, Tensor]]] = None,
|
||||
passages: Union[Dict[str, Tensor], List[Dict[str, Tensor]]] = None,
|
||||
teacher_scores: Union[None, List[float]] = None,
|
||||
no_in_batch_neg_flag: bool = False,
|
||||
):
|
||||
"""The computation performed at every call.
|
||||
|
||||
Args:
|
||||
queries (Union[Dict[str, Tensor], List[Dict[str, Tensor]]], optional): Input queries. Defaults to ``None``.
|
||||
passages (Union[Dict[str, Tensor], List[Dict[str, Tensor]]], optional): Input passages. Defaults to ``None``.
|
||||
teacher_scores (Union[None, List[float]], optional): Teacher scores for distillation. Defaults to ``None``.
|
||||
no_in_batch_neg_flag (bool, optional): If True, use no in-batch negatives and no cross-device negatives. Defaults to ``False``.
|
||||
|
||||
Returns:
|
||||
EmbedderOutput: Output of the forward call of model.
|
||||
"""
|
||||
q_reps = self.encode(queries) # (batch_size, dim)
|
||||
p_reps = self.encode(passages) # (batch_size * group_size, dim)
|
||||
|
||||
if self.use_mrl:
|
||||
device = q_reps[0].device
|
||||
batch_size = q_reps[0].size(0)
|
||||
else:
|
||||
device = q_reps.device
|
||||
batch_size = q_reps.size(0)
|
||||
|
||||
if self.training:
|
||||
if teacher_scores is not None:
|
||||
teacher_scores = torch.tensor(teacher_scores, device=device)
|
||||
teacher_scores = teacher_scores.view(batch_size, -1).detach() # (batch_size, group_size)
|
||||
teacher_targets = F.softmax(teacher_scores, dim=-1) # (batch_size, group_size)
|
||||
else:
|
||||
teacher_targets = None
|
||||
|
||||
if no_in_batch_neg_flag:
|
||||
compute_loss_func = self._compute_no_in_batch_neg_loss
|
||||
else:
|
||||
if self.negatives_cross_device:
|
||||
compute_loss_func = self._compute_cross_device_neg_loss
|
||||
else:
|
||||
compute_loss_func = self._compute_in_batch_neg_loss
|
||||
|
||||
if self.use_mrl:
|
||||
# compute MRL loss
|
||||
all_loss = torch.tensor(0.0, device=device)
|
||||
for dim_q_reps, dim_p_reps in zip(q_reps, p_reps):
|
||||
_, mrl_loss = compute_loss_func(dim_q_reps, dim_p_reps, teacher_targets=teacher_targets)
|
||||
all_loss += mrl_loss
|
||||
loss = all_loss / len(self.mrl_dims)
|
||||
else:
|
||||
scores, loss = compute_loss_func(q_reps, p_reps, teacher_targets=teacher_targets)
|
||||
else:
|
||||
loss = None
|
||||
|
||||
return EmbedderOutput(
|
||||
loss=loss,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def distill_loss(kd_loss_type, teacher_targets, student_scores, group_size=None):
|
||||
"""Compute the distillation loss.
|
||||
|
||||
Args:
|
||||
kd_loss_type (str): Type of knowledge distillation loss, supports "kl_div" and "m3_kd_loss".
|
||||
teacher_targets (torch.Tensor): Targets from the teacher model.
|
||||
student_scores (torch.Tensor): Score of student model.
|
||||
group_size (int, optional): Number of groups for . Defaults to ``None``.
|
||||
|
||||
Raises:
|
||||
ValueError: Invalid kd_loss_type
|
||||
|
||||
Returns:
|
||||
torch.Tensor: A scalar of computed distillation loss.
|
||||
"""
|
||||
if kd_loss_type == 'kl_div':
|
||||
# teacher_targets: (batch_size, group_size) / (world_size * batch_size, group_size)
|
||||
# student_scores: (batch_size, group_size) / (world_size * batch_size, group_size)
|
||||
return - torch.mean(
|
||||
torch.sum(torch.log_softmax(student_scores, dim=-1) * teacher_targets, dim=-1)
|
||||
)
|
||||
elif kd_loss_type == 'm3_kd_loss':
|
||||
# teacher_targets: (batch_size, group_size) / (world_size * batch_size, group_size)
|
||||
# student_scores: (batch_size, batch_size * group_size) / (world_size * batch_size, world_size * batch_size * group_size)
|
||||
labels = torch.arange(student_scores.size(0), device=student_scores.device, dtype=torch.long)
|
||||
labels = labels * group_size
|
||||
|
||||
loss = 0
|
||||
mask = torch.zeros_like(student_scores)
|
||||
for i in range(group_size):
|
||||
temp_target = labels + i
|
||||
temp_scores = student_scores + mask
|
||||
temp_loss = F.cross_entropy(temp_scores, temp_target, reduction="none") # B
|
||||
loss += torch.mean(teacher_targets[:, i] * temp_loss)
|
||||
mask = torch.scatter(mask, dim=-1, index=temp_target.unsqueeze(-1),
|
||||
value=torch.finfo(student_scores.dtype).min)
|
||||
return loss
|
||||
else:
|
||||
raise ValueError(f"Invalid kd_loss_type: {kd_loss_type}")
|
||||
|
||||
def _dist_gather_tensor(self, t: Optional[torch.Tensor]):
|
||||
"""Gather a tensor from all processes in a distributed setting.
|
||||
|
||||
Args:
|
||||
t (Optional[torch.Tensor]): The input tensor to be gathered. If `None`, no gathering is performed.
|
||||
|
||||
Returns:
|
||||
Union[torch.Tensor, None]: A concatenated tensor from all processes if ``t`` is not ``None``,
|
||||
otherwise returns ``None``.
|
||||
"""
|
||||
if t is None:
|
||||
return None
|
||||
t = t.contiguous()
|
||||
|
||||
all_tensors = [torch.empty_like(t) for _ in range(self.world_size)]
|
||||
dist.all_gather(all_tensors, t)
|
||||
|
||||
all_tensors[self.process_rank] = t
|
||||
all_tensors = torch.cat(all_tensors, dim=0)
|
||||
|
||||
return all_tensors
|
||||
@@ -0,0 +1,150 @@
|
||||
import os
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Tuple
|
||||
from abc import ABC, abstractmethod
|
||||
from transformers import set_seed, PreTrainedTokenizer
|
||||
|
||||
|
||||
from .AbsArguments import (
|
||||
AbsEmbedderModelArguments,
|
||||
AbsEmbedderDataArguments,
|
||||
AbsEmbedderTrainingArguments
|
||||
)
|
||||
from .AbsTrainer import AbsEmbedderTrainer
|
||||
from .AbsModeling import AbsEmbedderModel
|
||||
from .AbsDataset import (
|
||||
AbsEmbedderTrainDataset, AbsEmbedderCollator,
|
||||
AbsEmbedderSameDatasetTrainDataset, AbsEmbedderSameDatasetCollator
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsEmbedderRunner(ABC):
|
||||
"""Abstract class to run embedding model fine-tuning.
|
||||
|
||||
Args:
|
||||
model_args (AbsEmbedderModelArguments): Model arguments
|
||||
data_args (AbsEmbedderDataArguments): Data arguments.
|
||||
training_args (AbsEmbedderTrainingArguments): Training arguments.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
model_args: AbsEmbedderModelArguments,
|
||||
data_args: AbsEmbedderDataArguments,
|
||||
training_args: AbsEmbedderTrainingArguments
|
||||
):
|
||||
self.model_args = model_args
|
||||
self.data_args = data_args
|
||||
self.training_args = training_args
|
||||
|
||||
if (
|
||||
os.path.exists(training_args.output_dir)
|
||||
and os.listdir(training_args.output_dir)
|
||||
and training_args.do_train
|
||||
and not training_args.overwrite_output_dir
|
||||
):
|
||||
raise ValueError(
|
||||
f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome."
|
||||
)
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
|
||||
)
|
||||
logger.warning(
|
||||
"Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
|
||||
training_args.local_rank,
|
||||
training_args.device,
|
||||
training_args.n_gpu,
|
||||
bool(training_args.local_rank != -1),
|
||||
training_args.fp16,
|
||||
)
|
||||
logger.info("Training/evaluation parameters %s", training_args)
|
||||
logger.info("Model parameters %s", model_args)
|
||||
logger.info("Data parameters %s", data_args)
|
||||
|
||||
# Set seed
|
||||
set_seed(training_args.seed)
|
||||
|
||||
self.tokenizer, self.model = self.load_tokenizer_and_model()
|
||||
self.train_dataset = self.load_train_dataset()
|
||||
self.data_collator = self.load_data_collator()
|
||||
self.trainer = self.load_trainer()
|
||||
|
||||
@abstractmethod
|
||||
def load_tokenizer_and_model(self) -> Tuple[PreTrainedTokenizer, AbsEmbedderModel]:
|
||||
"""Abstract method to load the tokenizer and model.
|
||||
|
||||
Returns:
|
||||
Tuple[PreTrainedTokenizer, AbsEmbedderModel]: Loaded tokenizer and model instances.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def load_trainer(self) -> AbsEmbedderTrainer:
|
||||
"""Abstract method to load the trainer.
|
||||
|
||||
Returns:
|
||||
AbsEmbedderTrainer: The loaded trainer instance.
|
||||
"""
|
||||
pass
|
||||
|
||||
def load_train_dataset(self) -> AbsEmbedderTrainDataset:
|
||||
"""Loads the training dataset based on data arguments.
|
||||
|
||||
Returns:
|
||||
AbsEmbedderTrainDataset: The loaded dataset instance.
|
||||
"""
|
||||
if self.data_args.same_dataset_within_batch:
|
||||
train_dataset = AbsEmbedderSameDatasetTrainDataset(
|
||||
args=self.data_args,
|
||||
default_batch_size=self.training_args.per_device_train_batch_size,
|
||||
seed=self.training_args.seed,
|
||||
tokenizer=self.tokenizer,
|
||||
process_index=self.training_args.process_index,
|
||||
num_processes=self.training_args.world_size
|
||||
)
|
||||
self.training_args.per_device_train_batch_size = 1
|
||||
self.training_args.dataloader_num_workers = 0 # avoid multi-processing
|
||||
else:
|
||||
train_dataset = AbsEmbedderTrainDataset(
|
||||
args=self.data_args,
|
||||
tokenizer=self.tokenizer
|
||||
)
|
||||
return train_dataset
|
||||
|
||||
def load_data_collator(self) -> AbsEmbedderCollator:
|
||||
"""Loads the appropriate data collator.
|
||||
|
||||
Returns:
|
||||
AbsEmbedderCollator: Loaded data collator.
|
||||
"""
|
||||
if self.data_args.same_dataset_within_batch:
|
||||
EmbedCollator = AbsEmbedderSameDatasetCollator
|
||||
else:
|
||||
EmbedCollator = AbsEmbedderCollator
|
||||
|
||||
data_collator = EmbedCollator(
|
||||
tokenizer=self.tokenizer,
|
||||
query_max_len=self.data_args.query_max_len,
|
||||
passage_max_len=self.data_args.passage_max_len,
|
||||
sub_batch_size=self.training_args.sub_batch_size,
|
||||
pad_to_multiple_of=self.data_args.pad_to_multiple_of,
|
||||
padding=True,
|
||||
return_tensors="pt"
|
||||
)
|
||||
return data_collator
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
Executes the training process.
|
||||
"""
|
||||
Path(self.training_args.output_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Training
|
||||
self.trainer.train(resume_from_checkpoint=self.training_args.resume_from_checkpoint)
|
||||
self.trainer.save_model()
|
||||
@@ -0,0 +1,37 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
from abc import ABC, abstractmethod
|
||||
from transformers.trainer import Trainer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsEmbedderTrainer(ABC, Trainer):
|
||||
"""
|
||||
Abstract class for the trainer of embedder.
|
||||
"""
|
||||
@abstractmethod
|
||||
def _save(self, output_dir: Optional[str] = None, state_dict=None):
|
||||
pass
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
|
||||
"""
|
||||
How the loss is computed by Trainer. By default, all models return the loss in the first element.
|
||||
|
||||
Subclass and override for custom behavior.
|
||||
|
||||
Args:
|
||||
model (AbsEmbedderModel): The model being trained.
|
||||
inputs (dict): A dictionary of input tensors to be passed to the model.
|
||||
return_outputs (bool, optional): If ``True``, returns both the loss and the model's outputs. Otherwise,
|
||||
returns only the loss.
|
||||
|
||||
Returns:
|
||||
Union[torch.Tensor, tuple(torch.Tensor, EmbedderOutput)]: The computed loss. If ``return_outputs`` is ``True``,
|
||||
also returns the model's outputs in a tuple ``(loss, outputs)``.
|
||||
"""
|
||||
|
||||
outputs = model(**inputs)
|
||||
loss = outputs.loss
|
||||
|
||||
return (loss, outputs) if return_outputs else loss
|
||||
@@ -0,0 +1,30 @@
|
||||
from .AbsArguments import (
|
||||
AbsEmbedderDataArguments,
|
||||
AbsEmbedderModelArguments,
|
||||
AbsEmbedderTrainingArguments,
|
||||
)
|
||||
from .AbsDataset import (
|
||||
AbsEmbedderCollator, AbsEmbedderSameDatasetCollator,
|
||||
AbsEmbedderSameDatasetTrainDataset,
|
||||
AbsEmbedderTrainDataset,
|
||||
EmbedderTrainerCallbackForDataRefresh,
|
||||
)
|
||||
from .AbsModeling import AbsEmbedderModel, EmbedderOutput
|
||||
from .AbsTrainer import AbsEmbedderTrainer
|
||||
from .AbsRunner import AbsEmbedderRunner
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AbsEmbedderModelArguments",
|
||||
"AbsEmbedderDataArguments",
|
||||
"AbsEmbedderTrainingArguments",
|
||||
"AbsEmbedderModel",
|
||||
"AbsEmbedderTrainer",
|
||||
"AbsEmbedderRunner",
|
||||
"AbsEmbedderTrainDataset",
|
||||
"AbsEmbedderCollator",
|
||||
"AbsEmbedderSameDatasetTrainDataset",
|
||||
"AbsEmbedderSameDatasetCollator",
|
||||
"EmbedderOutput",
|
||||
"EmbedderTrainerCallbackForDataRefresh",
|
||||
]
|
||||
@@ -0,0 +1,141 @@
|
||||
import os
|
||||
from typing import Optional
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from transformers import TrainingArguments
|
||||
|
||||
|
||||
@dataclass
|
||||
class AbsRerankerModelArguments:
|
||||
"""
|
||||
Abstract class for reranker model arguments.
|
||||
"""
|
||||
|
||||
model_name_or_path: str = field(
|
||||
metadata={"help": "The model checkpoint for initialization."}
|
||||
)
|
||||
config_name: str = field(
|
||||
default=None,
|
||||
metadata={"help": "Pretrained config name or path if not the same as model_name."}
|
||||
)
|
||||
tokenizer_name: str = field(
|
||||
default=None,
|
||||
metadata={"help": "Pretrained tokenizer name or path if not the same as model_name."}
|
||||
)
|
||||
cache_dir: str = field(
|
||||
default=None,
|
||||
metadata={"help": "Where do you want to store the pre-trained models downloaded from s3."}
|
||||
)
|
||||
trust_remote_code: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "Trust remote code"}
|
||||
)
|
||||
model_type: str = field(
|
||||
default='encoder',
|
||||
metadata={"help": "Type of finetune, ['encoder', 'decoder']"}
|
||||
)
|
||||
use_fast_tokenizer: bool = field(
|
||||
default=True,
|
||||
metadata={"help": "Whether to use fast tokenizer or not."}
|
||||
)
|
||||
token: str = field(
|
||||
default_factory=lambda: os.getenv('HF_TOKEN', None),
|
||||
metadata={"help": "The token to use when accessing the model."}
|
||||
)
|
||||
# finetune_type: str = field(
|
||||
# default='sratch',
|
||||
# metadata={"help": "Type of finetune, ['sratch', 'finetune']"}
|
||||
# )
|
||||
|
||||
|
||||
@dataclass
|
||||
class AbsRerankerDataArguments:
|
||||
"""
|
||||
Abstract class for reranker data arguments.
|
||||
"""
|
||||
train_data: str = field(
|
||||
default=None, metadata={
|
||||
"help": "One or more paths to training data. `query: str`, `pos: List[str]`, `neg: List[str]` are required in the training data.",
|
||||
"nargs": "+"
|
||||
}
|
||||
)
|
||||
cache_path: Optional[str] = field(
|
||||
default=None, metadata={"help": "Where do you want to store the cached data"}
|
||||
)
|
||||
train_group_size: int = field(default=8)
|
||||
|
||||
query_max_len: int = field(
|
||||
default=32,
|
||||
metadata={
|
||||
"help": "The maximum total input sequence length after tokenization for passage. Sequences longer than this will be truncated."
|
||||
},
|
||||
)
|
||||
|
||||
passage_max_len: int = field(
|
||||
default=128,
|
||||
metadata={
|
||||
"help": "The maximum total input sequence length after tokenization for passage. Sequences longer than this will be truncated."
|
||||
},
|
||||
)
|
||||
|
||||
max_len: int = field(
|
||||
default=512,
|
||||
metadata={
|
||||
"help": "The maximum total input sequence length after tokenization. Sequences longer than this will be truncated."
|
||||
},
|
||||
)
|
||||
|
||||
pad_to_multiple_of: Optional[int] = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": "If set will pad the sequence to be a multiple of the provided value."
|
||||
},
|
||||
)
|
||||
|
||||
max_example_num_per_dataset: int = field(
|
||||
default=100000000, metadata={"help": "the max number of examples for each dataset"}
|
||||
)
|
||||
|
||||
query_instruction_for_rerank: str= field(
|
||||
default=None, metadata={"help": "instruction for query"}
|
||||
)
|
||||
query_instruction_format: str = field(
|
||||
default="{}{}", metadata={"help": "format for query instruction"}
|
||||
)
|
||||
|
||||
knowledge_distillation: bool = field(
|
||||
default=False,
|
||||
metadata={"help": "Use knowledge distillation when `pos_scores: List[float]` and `neg_scores: List[float]` are in features of training data"}
|
||||
)
|
||||
|
||||
passage_instruction_for_rerank: Optional[str] = field(
|
||||
default=None, metadata={"help": "instruction for passage"}
|
||||
)
|
||||
passage_instruction_format: Optional[str] = field(
|
||||
default="{}{}", metadata={"help": "format for passage instruction"}
|
||||
)
|
||||
|
||||
shuffle_ratio: float = field(
|
||||
default=0.0, metadata={"help": "The ratio of shuffling the text"}
|
||||
)
|
||||
|
||||
sep_token: str = field(
|
||||
default='\n', metadata={"help": "The sep token for LLM reranker to discriminate between query and passage"}
|
||||
)
|
||||
|
||||
def __post_init__(self):
|
||||
# replace "\\n" with "\n"
|
||||
if "\\n" in self.query_instruction_format:
|
||||
self.query_instruction_format = self.query_instruction_format.replace("\\n", "\n")
|
||||
if "\\n" in self.passage_instruction_format:
|
||||
self.passage_instruction_format = self.passage_instruction_format.replace("\\n", "\n")
|
||||
|
||||
# check the existence of train data
|
||||
for train_dir in self.train_data:
|
||||
if not os.path.exists(train_dir):
|
||||
raise FileNotFoundError(f"cannot find file: {train_dir}, please set a true path")
|
||||
|
||||
|
||||
@dataclass
|
||||
class AbsRerankerTrainingArguments(TrainingArguments):
|
||||
sub_batch_size: Optional[int] = field(default=None, metadata={"help": "sub batch size for training, not implemented yet"})
|
||||
@@ -0,0 +1,401 @@
|
||||
import os
|
||||
import math
|
||||
import random
|
||||
import logging
|
||||
import datasets
|
||||
import numpy as np
|
||||
import torch.distributed as dist
|
||||
from dataclasses import dataclass
|
||||
from torch.utils.data import Dataset
|
||||
from transformers import (
|
||||
PreTrainedTokenizer,
|
||||
DataCollatorWithPadding,
|
||||
BatchEncoding,
|
||||
DataCollatorForSeq2Seq
|
||||
)
|
||||
from typing import List
|
||||
|
||||
from .AbsArguments import AbsRerankerDataArguments
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsRerankerTrainDataset(Dataset):
|
||||
"""Abstract class for reranker training dataset.
|
||||
|
||||
Args:
|
||||
args (AbsRerankerDataArguments): Data arguments.
|
||||
tokenizer (PreTrainedTokenizer): Tokenizer to use.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
args: AbsRerankerDataArguments,
|
||||
tokenizer: PreTrainedTokenizer
|
||||
):
|
||||
self.args = args
|
||||
self.tokenizer = tokenizer
|
||||
|
||||
train_datasets = []
|
||||
for data_dir in args.train_data:
|
||||
if not os.path.isdir(data_dir):
|
||||
if not (data_dir.endswith('.json') or data_dir.endswith('.jsonl')): continue
|
||||
temp_dataset = self._load_dataset(data_dir)
|
||||
if len(temp_dataset) == 0: continue
|
||||
train_datasets.append(temp_dataset)
|
||||
else:
|
||||
for file in os.listdir(data_dir):
|
||||
if not (file.endswith('.json') or file.endswith('.jsonl')): continue
|
||||
temp_dataset = self._load_dataset(os.path.join(data_dir, file))
|
||||
if len(temp_dataset) == 0: continue
|
||||
train_datasets.append(temp_dataset)
|
||||
self.dataset = datasets.concatenate_datasets(train_datasets)
|
||||
|
||||
self.max_length = self.args.query_max_len + self.args.passage_max_len
|
||||
|
||||
def _load_dataset(self, file_path: str):
|
||||
"""Load dataset from path.
|
||||
|
||||
Args:
|
||||
file_path (str): Path to load the datasets from.
|
||||
|
||||
Raises:
|
||||
ValueError: `pos_scores` and `neg_scores` not found in the features of training data
|
||||
|
||||
Returns:
|
||||
datasets.Dataset: Loaded HF dataset.
|
||||
"""
|
||||
safe_rank = dist.get_rank() if dist.is_initialized() else 0
|
||||
if safe_rank == 0:
|
||||
logger.info(f'loading data from {file_path} ...')
|
||||
|
||||
temp_dataset = datasets.load_dataset('json', data_files=file_path, split='train', cache_dir=self.args.cache_path)
|
||||
if len(temp_dataset) > self.args.max_example_num_per_dataset:
|
||||
temp_dataset = temp_dataset.select(random.sample(list(range(len(temp_dataset))), self.args.max_example_num_per_dataset))
|
||||
if not self.args.knowledge_distillation:
|
||||
if 'pos_scores' in temp_dataset.column_names:
|
||||
temp_dataset = temp_dataset.remove_columns(['pos_scores'])
|
||||
if 'neg_scores' in temp_dataset.column_names:
|
||||
temp_dataset = temp_dataset.remove_columns(['neg_scores'])
|
||||
else:
|
||||
if 'pos_scores' not in temp_dataset.column_names or 'neg_scores' not in temp_dataset.column_names:
|
||||
raise ValueError(f"`pos_scores` and `neg_scores` not found in the features of training data in {file_path}, which is necessary when using knowledge distillation.")
|
||||
return temp_dataset
|
||||
|
||||
def _shuffle_text(self, text):
|
||||
"""shuffle the input text.
|
||||
|
||||
Args:
|
||||
text (str): Input text.
|
||||
|
||||
Returns:
|
||||
str: Shuffled text.
|
||||
"""
|
||||
if self.args.shuffle_ratio > 0 and len(text) > 100 and random.random() < self.args.shuffle_ratio:
|
||||
split_text = []
|
||||
chunk_size = len(text)//3 + 1
|
||||
for i in range(0, len(text), chunk_size):
|
||||
split_text.append(text[i:i+chunk_size])
|
||||
random.shuffle(split_text)
|
||||
return " ".join(split_text)
|
||||
else:
|
||||
return text
|
||||
|
||||
def __len__(self):
|
||||
return len(self.dataset)
|
||||
|
||||
def create_one_example(self, qry_encoding: str, doc_encoding: str):
|
||||
"""Creates a single input example by encoding and preparing a query and document pair for the model.
|
||||
|
||||
Args:
|
||||
qry_encoding (str): Query to be encoded.
|
||||
doc_encoding (str): Document to be encoded.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing tokenized and prepared inputs, ready for model consumption.
|
||||
"""
|
||||
qry_inputs = self.tokenizer.encode(qry_encoding, truncation=True, max_length=self.args.query_max_len + self.args.passage_max_len // 4, add_special_tokens=False)
|
||||
doc_inputs = self.tokenizer.encode(doc_encoding, truncation=True, max_length=self.args.passage_max_len + self.args.query_max_len // 2, add_special_tokens=False)
|
||||
item = self.tokenizer.prepare_for_model(
|
||||
qry_inputs,
|
||||
doc_inputs,
|
||||
truncation='only_second',
|
||||
max_length=self.args.query_max_len + self.args.passage_max_len,
|
||||
padding=False,
|
||||
)
|
||||
return item
|
||||
|
||||
def __getitem__(self, item):
|
||||
data = self.dataset[item]
|
||||
train_group_size = self.args.train_group_size
|
||||
|
||||
query = data['query']
|
||||
if self.args.query_instruction_for_rerank is not None:
|
||||
query = self.args.query_instruction_format.format(
|
||||
data['query_prompt'] if 'query_prompt' in data else self.args.query_instruction_for_rerank,
|
||||
query
|
||||
)
|
||||
|
||||
passages = []
|
||||
teacher_scores = []
|
||||
|
||||
assert isinstance(data['pos'], list) and isinstance(data['neg'], list)
|
||||
|
||||
pos_idx = random.choice(list(range(len(data['pos']))))
|
||||
passages.append(self._shuffle_text(data['pos'][pos_idx]))
|
||||
|
||||
neg_all_idx = list(range(len(data['neg'])))
|
||||
if len(data['neg']) < train_group_size - 1:
|
||||
num = math.ceil((train_group_size - 1) / len(data['neg']))
|
||||
neg_idxs = random.sample(neg_all_idx * num, train_group_size - 1)
|
||||
else:
|
||||
neg_idxs = random.sample(neg_all_idx, self.args.train_group_size - 1)
|
||||
for neg_idx in neg_idxs:
|
||||
passages.append(data['neg'][neg_idx])
|
||||
|
||||
if self.args.knowledge_distillation:
|
||||
assert isinstance(data['pos_scores'], list) and isinstance(data['neg_scores'], list)
|
||||
teacher_scores.append(data['pos_scores'][pos_idx])
|
||||
for neg_idx in neg_idxs:
|
||||
teacher_scores.append(data['neg_scores'][neg_idx])
|
||||
if not all(isinstance(score, (int, float)) for score in teacher_scores):
|
||||
raise ValueError(f"pos_score or neg_score must be digit")
|
||||
else:
|
||||
teacher_scores = None
|
||||
|
||||
if self.args.passage_instruction_for_rerank is not None:
|
||||
passages = [
|
||||
self.args.passage_instruction_format.format(
|
||||
data['passage_prompt'] if 'passage_prompt' in data else self.args.passage_instruction_for_rerank, p
|
||||
)
|
||||
for p in passages
|
||||
]
|
||||
|
||||
batch_data = []
|
||||
for passage in passages:
|
||||
batch_data.append(self.create_one_example(query, passage))
|
||||
|
||||
return batch_data, teacher_scores
|
||||
|
||||
@dataclass
|
||||
class AbsRerankerCollator(DataCollatorWithPadding):
|
||||
"""
|
||||
The abstract reranker collator.
|
||||
"""
|
||||
query_max_len: int = 32
|
||||
passage_max_len: int = 128
|
||||
|
||||
def __call__(self, features) -> List[BatchEncoding]:
|
||||
teacher_scores = [f[1] for f in features]
|
||||
if teacher_scores[0] is None:
|
||||
teacher_scores = None
|
||||
elif isinstance(teacher_scores[0], list):
|
||||
teacher_scores = sum(teacher_scores, [])
|
||||
|
||||
features = [f[0] for f in features]
|
||||
if isinstance(features[0], list):
|
||||
features = sum(features, [])
|
||||
|
||||
collated = self.tokenizer.pad(
|
||||
features,
|
||||
padding=self.padding,
|
||||
max_length=self.query_max_len + self.passage_max_len,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
return_tensors=self.return_tensors,
|
||||
)
|
||||
|
||||
return {
|
||||
"pair": collated,
|
||||
"teacher_scores": teacher_scores,
|
||||
}
|
||||
|
||||
class AbsLLMRerankerTrainDataset(AbsRerankerTrainDataset):
|
||||
"""Abstract class for LLM reranker training dataset.
|
||||
|
||||
Args:
|
||||
args (AbsRerankerDataArguments): Data arguments.
|
||||
tokenizer (PreTrainedTokenizer): Tokenizer to use.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
args: AbsRerankerDataArguments,
|
||||
tokenizer: PreTrainedTokenizer
|
||||
):
|
||||
super().__init__(args, tokenizer)
|
||||
sep = self.args.sep_token
|
||||
self.sep_inputs = self.tokenizer(
|
||||
sep,
|
||||
return_tensors=None,
|
||||
add_special_tokens=False
|
||||
)['input_ids']
|
||||
|
||||
def __getitem__(self, item) -> List[BatchEncoding]:
|
||||
data = self.dataset[item]
|
||||
train_group_size = self.args.train_group_size
|
||||
|
||||
query = data['query']
|
||||
if self.args.query_instruction_for_rerank is not None:
|
||||
query = self.args.query_instruction_format.format(
|
||||
data['query_prompt'] if 'query_prompt' in data else self.args.query_instruction_for_rerank,
|
||||
query
|
||||
)
|
||||
|
||||
passages = []
|
||||
teacher_scores = []
|
||||
|
||||
assert isinstance(data['pos'], list) and isinstance(data['neg'], list)
|
||||
|
||||
pos_idx = random.choice(list(range(len(data['pos']))))
|
||||
passages.append(self._shuffle_text(data['pos'][pos_idx]))
|
||||
|
||||
neg_all_idx = list(range(len(data['neg'])))
|
||||
if len(data['neg']) < train_group_size - 1:
|
||||
num = math.ceil((train_group_size - 1) / len(data['neg']))
|
||||
neg_idxs = random.sample(neg_all_idx * num, train_group_size - 1)
|
||||
else:
|
||||
neg_idxs = random.sample(neg_all_idx, self.args.train_group_size - 1)
|
||||
for neg_idx in neg_idxs:
|
||||
passages.append(data['neg'][neg_idx])
|
||||
|
||||
if self.args.knowledge_distillation:
|
||||
assert isinstance(data['pos_scores'], list) and isinstance(data['neg_scores'], list)
|
||||
teacher_scores.append(data['pos_scores'][pos_idx])
|
||||
for neg_idx in neg_idxs:
|
||||
teacher_scores.append(data['neg_scores'][neg_idx])
|
||||
if not all(isinstance(score, (int, float)) for score in teacher_scores):
|
||||
raise ValueError(f"pos_score or neg_score must be digit")
|
||||
else:
|
||||
teacher_scores = None
|
||||
|
||||
if self.args.passage_instruction_for_rerank is not None:
|
||||
passages = [
|
||||
self.args.passage_instruction_format.format(
|
||||
data['passage_prompt'] if 'passage_prompt' in data else self.args.passage_instruction_for_rerank, p
|
||||
)
|
||||
for p in passages
|
||||
]
|
||||
|
||||
prompt = self.dataset[item].get('prompt', "Given a query A and a passage B, determine whether the passage contains an answer to the query by providing a prediction of either 'Yes' or 'No'.")
|
||||
|
||||
query_inputs = self.tokenizer(
|
||||
query,
|
||||
return_tensors=None,
|
||||
max_length=self.args.query_max_len + self.args.passage_max_len // 4,
|
||||
truncation=True,
|
||||
add_special_tokens=False
|
||||
)
|
||||
|
||||
prompt_inputs = self.tokenizer(
|
||||
prompt,
|
||||
return_tensors=None,
|
||||
add_special_tokens=False
|
||||
)['input_ids']
|
||||
|
||||
max_length = self.max_length - len(prompt_inputs) - len(self.sep_inputs)
|
||||
|
||||
passages_inputs = []
|
||||
for i, passage in enumerate(passages):
|
||||
passage_inputs = self.tokenizer(
|
||||
passage,
|
||||
return_tensors=None,
|
||||
max_length=self.args.passage_max_len + self.args.query_max_len // 2,
|
||||
truncation=True,
|
||||
add_special_tokens=False
|
||||
)
|
||||
if self.tokenizer.bos_token_id is not None and self.tokenizer.bos_token_id != self.tokenizer.pad_token_id:
|
||||
item = self.tokenizer.prepare_for_model(
|
||||
[self.tokenizer.bos_token_id] + query_inputs['input_ids'],
|
||||
self.sep_inputs + passage_inputs['input_ids'],
|
||||
truncation='only_second',
|
||||
max_length=max_length,
|
||||
padding=False,
|
||||
return_attention_mask=False,
|
||||
return_token_type_ids=False,
|
||||
add_special_tokens=False
|
||||
)
|
||||
else:
|
||||
item = self.tokenizer.prepare_for_model(
|
||||
query_inputs['input_ids'],
|
||||
self.sep_inputs + passage_inputs['input_ids'],
|
||||
truncation='only_second',
|
||||
max_length=max_length,
|
||||
padding=False,
|
||||
return_attention_mask=False,
|
||||
return_token_type_ids=False,
|
||||
add_special_tokens=False
|
||||
)
|
||||
|
||||
passage_inputs['input_ids'] = item['input_ids'] + self.sep_inputs + prompt_inputs
|
||||
|
||||
passage_inputs['attention_mask'] = [1] * len(passage_inputs['input_ids'])
|
||||
# passage_inputs['labels'] = passage_inputs['input_ids'].copy()
|
||||
# passage_inputs['labels'] = [-100] * (len(passage_inputs['input_ids']) - 1) + passage_inputs['labels'][(len(passage_inputs['input_ids']) - 1):]
|
||||
passage_inputs.pop('token_type_ids') if 'token_type_ids' in passage_inputs.keys() else None
|
||||
if 'position_ids' in passage_inputs.keys():
|
||||
passage_inputs['position_ids'] = list(range(len(passage_inputs['input_ids'])))
|
||||
passages_inputs.append(passage_inputs)
|
||||
|
||||
return passages_inputs, teacher_scores
|
||||
|
||||
|
||||
@dataclass
|
||||
class AbsLLMRerankerCollator(DataCollatorForSeq2Seq):
|
||||
"""
|
||||
Wrapper that does conversion from List[Tuple[encode_qry, encode_psg]] to List[qry], List[psg]
|
||||
and pass batch separately to the actual collator.
|
||||
Abstract out data detail for the model.
|
||||
"""
|
||||
query_max_len: int = 32
|
||||
passage_max_len: int = 128
|
||||
|
||||
def __call__(self, features, return_tensors='pt'):
|
||||
if return_tensors is None:
|
||||
return_tensors = self.return_tensors
|
||||
|
||||
teacher_scores = [f[1] for f in features]
|
||||
if teacher_scores[0] is None:
|
||||
teacher_scores = None
|
||||
elif isinstance(teacher_scores[0], list):
|
||||
teacher_scores = sum(teacher_scores, [])
|
||||
|
||||
features = [f[0] for f in features]
|
||||
if isinstance(features[0], list):
|
||||
features = sum(features, [])
|
||||
|
||||
labels = [feature["labels"] for feature in features] if "labels" in features[0].keys() else None
|
||||
# We have to pad the labels before calling `tokenizer.pad` as this method won't pad them and needs them of the
|
||||
# same length to return tensors.
|
||||
if labels is not None:
|
||||
max_label_length = max(len(l) for l in labels)
|
||||
# print(max_label_length)
|
||||
if self.pad_to_multiple_of is not None:
|
||||
max_label_length = (
|
||||
(max_label_length + self.pad_to_multiple_of - 1)
|
||||
// self.pad_to_multiple_of
|
||||
* self.pad_to_multiple_of
|
||||
)
|
||||
|
||||
padding_side = self.tokenizer.padding_side
|
||||
for feature in features:
|
||||
remainder = [self.label_pad_token_id] * (max_label_length - len(feature["labels"]))
|
||||
if isinstance(feature["labels"], list):
|
||||
feature["labels"] = (
|
||||
feature["labels"] + remainder
|
||||
if padding_side == "right" else remainder + feature["labels"]
|
||||
)
|
||||
elif padding_side == "right":
|
||||
feature["labels"] = np.concatenate([feature["labels"], remainder]).astype(np.int64)
|
||||
else:
|
||||
feature["labels"] = np.concatenate([remainder, feature["labels"]]).astype(np.int64)
|
||||
|
||||
collated = self.tokenizer.pad(
|
||||
features,
|
||||
padding=self.padding,
|
||||
max_length=self.query_max_len + self.passage_max_len,
|
||||
return_tensors=return_tensors,
|
||||
pad_to_multiple_of=self.pad_to_multiple_of,
|
||||
)
|
||||
|
||||
return {
|
||||
"pair": collated,
|
||||
"teacher_scores": teacher_scores,
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
import torch
|
||||
from torch import nn, Tensor
|
||||
from transformers import PreTrainedTokenizer
|
||||
from transformers.file_utils import ModelOutput
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Dict, Optional, List, Union
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RerankerOutput(ModelOutput):
|
||||
loss: Optional[Tensor] = None
|
||||
scores: Optional[Tensor] = None
|
||||
|
||||
|
||||
class AbsRerankerModel(ABC, nn.Module):
|
||||
"""Abstract class of embedding model for training.
|
||||
|
||||
Args:
|
||||
base_model: The base model to train on.
|
||||
tokenizer (PreTrainedTokenizer, optional): The tokenizer to use. Defaults to ``None``.
|
||||
train_batch_size (int, optional): Batch size used for training. Defaults to ``4``.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
base_model: None,
|
||||
tokenizer: PreTrainedTokenizer = None,
|
||||
train_batch_size: int = 4,
|
||||
):
|
||||
nn.Module.__init__(self)
|
||||
self.model = base_model
|
||||
self.tokenizer = tokenizer
|
||||
self.cross_entropy = nn.CrossEntropyLoss(reduction='mean')
|
||||
|
||||
if self.model.config.pad_token_id is None:
|
||||
self.model.config.pad_token_id = self.tokenizer.pad_token_id
|
||||
self.config = self.model.config
|
||||
|
||||
self.train_batch_size = train_batch_size
|
||||
|
||||
self.yes_loc = self.tokenizer('Yes', add_special_tokens=False)['input_ids'][-1]
|
||||
|
||||
def gradient_checkpointing_enable(self, **kwargs):
|
||||
"""
|
||||
Activates gradient checkpointing for the current model.
|
||||
"""
|
||||
self.model.gradient_checkpointing_enable(**kwargs)
|
||||
|
||||
def enable_input_require_grads(self, **kwargs):
|
||||
"""
|
||||
Enables the gradients for the input embeddings.
|
||||
"""
|
||||
self.model.enable_input_require_grads(**kwargs)
|
||||
|
||||
@abstractmethod
|
||||
def encode(self, features):
|
||||
"""Abstract method of encode.
|
||||
|
||||
Args:
|
||||
features (dict): Teatures to pass to the model.
|
||||
"""
|
||||
pass
|
||||
|
||||
def forward(self, pair: Union[Dict[str, Tensor], List[Dict[str, Tensor]]] = None, teacher_scores: Optional[Tensor] = None):
|
||||
"""The computation performed at every call.
|
||||
|
||||
Args:
|
||||
pair (Union[Dict[str, Tensor], List[Dict[str, Tensor]]], optional): The query-document pair. Defaults to ``None``.
|
||||
teacher_scores (Optional[Tensor], optional): Teacher scores of knowledge distillation. Defaults to None.
|
||||
|
||||
Returns:
|
||||
RerankerOutput: Output of reranker model.
|
||||
"""
|
||||
ranker_logits = self.encode(pair) # (batch_size * num, dim)
|
||||
if teacher_scores is not None:
|
||||
teacher_scores = torch.Tensor(teacher_scores)
|
||||
teacher_targets = teacher_scores.view(self.train_batch_size, -1)
|
||||
teacher_targets = torch.softmax(teacher_targets.detach(), dim=-1)
|
||||
|
||||
if self.training:
|
||||
grouped_logits = ranker_logits.view(self.train_batch_size, -1)
|
||||
target = torch.zeros(self.train_batch_size, device=grouped_logits.device, dtype=torch.long)
|
||||
loss = self.compute_loss(grouped_logits, target)
|
||||
if teacher_scores is not None:
|
||||
teacher_targets = teacher_targets.to(grouped_logits.device)
|
||||
# print(teacher_targets, torch.mean(torch.sum(torch.log_softmax(grouped_logits, dim=-1) * teacher_targets, dim=-1)))
|
||||
loss += - torch.mean(torch.sum(torch.log_softmax(grouped_logits, dim=-1) * teacher_targets, dim=-1))
|
||||
else:
|
||||
loss = None
|
||||
|
||||
# print(loss)
|
||||
return RerankerOutput(
|
||||
loss=loss,
|
||||
scores=ranker_logits,
|
||||
)
|
||||
|
||||
def compute_loss(self, scores, target):
|
||||
"""Compute the loss.
|
||||
|
||||
Args:
|
||||
scores (torch.Tensor): Computed scores.
|
||||
target (torch.Tensor): The target value.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The computed loss.
|
||||
"""
|
||||
return self.cross_entropy(scores, target)
|
||||
|
||||
def save(self, output_dir: str):
|
||||
"""Save the model.
|
||||
|
||||
Args:
|
||||
output_dir (str): Directory for saving the model.
|
||||
"""
|
||||
# self.model.save_pretrained(output_dir)
|
||||
state_dict = self.model.state_dict()
|
||||
state_dict = type(state_dict)(
|
||||
{k: v.clone().cpu()
|
||||
for k,
|
||||
v in state_dict.items()})
|
||||
self.model.save_pretrained(output_dir, state_dict=state_dict)
|
||||
|
||||
def save_pretrained(self, *args, **kwargs):
|
||||
"""
|
||||
Save the tokenizer and model.
|
||||
"""
|
||||
self.tokenizer.save_pretrained(*args, **kwargs)
|
||||
return self.model.save_pretrained(*args, **kwargs)
|
||||
@@ -0,0 +1,143 @@
|
||||
import os
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Tuple
|
||||
from abc import ABC, abstractmethod
|
||||
from transformers import set_seed, PreTrainedTokenizer
|
||||
|
||||
|
||||
from .AbsArguments import (
|
||||
AbsRerankerModelArguments,
|
||||
AbsRerankerDataArguments,
|
||||
AbsRerankerTrainingArguments
|
||||
)
|
||||
from .AbsTrainer import AbsRerankerTrainer
|
||||
from .AbsModeling import AbsRerankerModel
|
||||
from .AbsDataset import (
|
||||
AbsRerankerTrainDataset, AbsRerankerCollator,
|
||||
AbsLLMRerankerTrainDataset, AbsLLMRerankerCollator
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsRerankerRunner(ABC):
|
||||
"""Abstract class to run reranker model fine-tuning.
|
||||
|
||||
Args:
|
||||
model_args (AbsRerankerModelArguments): Model arguments
|
||||
data_args (AbsRerankerDataArguments): Data arguments.
|
||||
training_args (AbsRerankerTrainingArguments): Training arguments.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
model_args: AbsRerankerModelArguments,
|
||||
data_args: AbsRerankerDataArguments,
|
||||
training_args: AbsRerankerTrainingArguments
|
||||
):
|
||||
self.model_args = model_args
|
||||
self.data_args = data_args
|
||||
self.training_args = training_args
|
||||
|
||||
if (
|
||||
os.path.exists(training_args.output_dir)
|
||||
and os.listdir(training_args.output_dir)
|
||||
and training_args.do_train
|
||||
and not training_args.overwrite_output_dir
|
||||
):
|
||||
raise ValueError(
|
||||
f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome."
|
||||
)
|
||||
|
||||
# Setup logging
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
|
||||
)
|
||||
logger.warning(
|
||||
"Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
|
||||
training_args.local_rank,
|
||||
training_args.device,
|
||||
training_args.n_gpu,
|
||||
bool(training_args.local_rank != -1),
|
||||
training_args.fp16,
|
||||
)
|
||||
logger.info("Training/evaluation parameters %s", training_args)
|
||||
logger.info("Model parameters %s", model_args)
|
||||
logger.info("Data parameters %s", data_args)
|
||||
|
||||
# Set seed
|
||||
set_seed(training_args.seed)
|
||||
|
||||
self.tokenizer, self.model = self.load_tokenizer_and_model()
|
||||
self.train_dataset = self.load_train_dataset()
|
||||
self.data_collator = self.load_data_collator()
|
||||
self.trainer = self.load_trainer()
|
||||
|
||||
@abstractmethod
|
||||
def load_tokenizer_and_model(self) -> Tuple[PreTrainedTokenizer, AbsRerankerModel]:
|
||||
"""Abstract method to load the tokenizer and model.
|
||||
|
||||
Returns:
|
||||
Tuple[PreTrainedTokenizer, AbsRerankerModel]: Loaded tokenizer and model instances.
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def load_trainer(self) -> AbsRerankerTrainer:
|
||||
"""Abstract method to load the trainer.
|
||||
|
||||
Returns:
|
||||
AbsRerankerTrainer: The loaded trainer instance.
|
||||
"""
|
||||
pass
|
||||
|
||||
def load_train_dataset(self) -> AbsRerankerTrainDataset:
|
||||
"""Loads the training dataset based on data arguments.
|
||||
|
||||
Returns:
|
||||
AbsRerankerTrainDataset: The loaded dataset instance.
|
||||
"""
|
||||
if self.model_args.model_type == 'encoder':
|
||||
train_dataset = AbsRerankerTrainDataset(
|
||||
args=self.data_args,
|
||||
tokenizer=self.tokenizer
|
||||
)
|
||||
else:
|
||||
train_dataset = AbsLLMRerankerTrainDataset(
|
||||
args=self.data_args,
|
||||
tokenizer=self.tokenizer
|
||||
)
|
||||
return train_dataset
|
||||
|
||||
def load_data_collator(self) -> AbsRerankerCollator:
|
||||
"""Loads the appropriate data collator.
|
||||
|
||||
Returns:
|
||||
AbsRerankerCollator: Loaded data collator.
|
||||
"""
|
||||
if self.model_args.model_type == 'encoder':
|
||||
RerankerCollator = AbsRerankerCollator
|
||||
else:
|
||||
RerankerCollator = AbsLLMRerankerCollator
|
||||
|
||||
data_collator = RerankerCollator(
|
||||
tokenizer=self.tokenizer,
|
||||
query_max_len=self.data_args.query_max_len,
|
||||
passage_max_len=self.data_args.passage_max_len,
|
||||
pad_to_multiple_of=self.data_args.pad_to_multiple_of,
|
||||
padding=True,
|
||||
return_tensors="pt"
|
||||
)
|
||||
return data_collator
|
||||
|
||||
def run(self):
|
||||
"""
|
||||
Executes the training process.
|
||||
"""
|
||||
Path(self.training_args.output_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Training
|
||||
self.trainer.train(resume_from_checkpoint=self.training_args.resume_from_checkpoint)
|
||||
self.trainer.save_model()
|
||||
@@ -0,0 +1,37 @@
|
||||
import logging
|
||||
from typing import Optional
|
||||
from abc import ABC, abstractmethod
|
||||
from transformers.trainer import Trainer
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsRerankerTrainer(ABC, Trainer):
|
||||
"""
|
||||
Abstract class for the trainer of reranker.
|
||||
"""
|
||||
@abstractmethod
|
||||
def _save(self, output_dir: Optional[str] = None, state_dict=None):
|
||||
pass
|
||||
|
||||
def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
|
||||
"""
|
||||
How the loss is computed by Trainer. By default, all models return the loss in the first element.
|
||||
|
||||
Subclass and override for custom behavior.
|
||||
|
||||
Args:
|
||||
model (AbsRerankerModel): The model being trained.
|
||||
inputs (dict): A dictionary of input tensors to be passed to the model.
|
||||
return_outputs (bool, optional): If ``True``, returns both the loss and the model's outputs. Otherwise,
|
||||
returns only the loss. Defaults to ``False``.
|
||||
|
||||
Returns:
|
||||
Union[torch.Tensor, tuple(torch.Tensor, RerankerOutput)]: The computed loss. If ``return_outputs`` is ``True``,
|
||||
also returns the model's outputs in a tuple ``(loss, outputs)``.
|
||||
"""
|
||||
|
||||
outputs = model(**inputs)
|
||||
loss = outputs.loss
|
||||
|
||||
return (loss, outputs) if return_outputs else loss
|
||||
@@ -0,0 +1,22 @@
|
||||
from .AbsArguments import AbsRerankerDataArguments, AbsRerankerModelArguments, AbsRerankerTrainingArguments
|
||||
from .AbsDataset import (
|
||||
AbsRerankerTrainDataset, AbsRerankerCollator,
|
||||
AbsLLMRerankerTrainDataset, AbsLLMRerankerCollator
|
||||
)
|
||||
from .AbsModeling import AbsRerankerModel, RerankerOutput
|
||||
from .AbsTrainer import AbsRerankerTrainer
|
||||
from .AbsRunner import AbsRerankerRunner
|
||||
|
||||
__all__ = [
|
||||
"AbsRerankerDataArguments",
|
||||
"AbsRerankerModelArguments",
|
||||
"AbsRerankerTrainingArguments",
|
||||
"AbsRerankerTrainDataset",
|
||||
"AbsRerankerCollator",
|
||||
"AbsLLMRerankerTrainDataset",
|
||||
"AbsLLMRerankerCollator",
|
||||
"AbsRerankerModel",
|
||||
"RerankerOutput",
|
||||
"AbsRerankerTrainer",
|
||||
"AbsRerankerRunner",
|
||||
]
|
||||
@@ -0,0 +1,491 @@
|
||||
import logging
|
||||
from tqdm import tqdm, trange
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Union, List, Dict, Literal, Optional
|
||||
|
||||
import queue
|
||||
import multiprocessing as mp
|
||||
from multiprocessing import Queue
|
||||
|
||||
import math
|
||||
import gc
|
||||
import torch
|
||||
import numpy as np
|
||||
from transformers import is_torch_npu_available
|
||||
|
||||
try:
|
||||
import torch_musa
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsEmbedder(ABC):
|
||||
"""
|
||||
Base class for embedder.
|
||||
Extend this class and implement :meth:`encode_queries`, :meth:`encode_corpus`, :meth:`encode` for custom embedders.
|
||||
|
||||
Args:
|
||||
model_name_or_path (str): If it's a path to a local model, it loads the model from the path. Otherwise tries to download and
|
||||
load a model from HuggingFace Hub with the name.
|
||||
normalize_embeddings (bool, optional): If True, normalize the embedding vector. Defaults to :data:`True`.
|
||||
use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance
|
||||
degradation. Defaults to :data:`True`.
|
||||
query_instruction_for_retrieval: (Optional[str], optional): Query instruction for retrieval tasks, which will be used with
|
||||
with :attr:`query_instruction_format`. Defaults to :data:`None`.
|
||||
query_instruction_format: (str, optional): The template for :attr:`query_instruction_for_retrieval`. Defaults to :data:`"{}{}"`.
|
||||
devices (Optional[Union[str, int, List[str], List[int]]], optional): Devices to use for model inference. Defaults to :data:`None`.
|
||||
batch_size (int, optional): Batch size for inference. Defaults to :data:`256`.
|
||||
query_max_length (int, optional): Maximum length for query. Defaults to :data:`512`.
|
||||
passage_max_length (int, optional): Maximum length for passage. Defaults to :data:`512`.
|
||||
convert_to_numpy (bool, optional): If True, the output embedding will be a Numpy array. Otherwise, it will be a Torch Tensor.
|
||||
Defaults to :data:`True`.
|
||||
truncate_dim (Optional[int], optional): The dimension to truncate the output embeddings to. Useful for Matryoshka
|
||||
Representation Learning models. If None, no truncation is performed. Defaults to :data:`None`.
|
||||
kwargs (Dict[Any], optional): Additional parameters for HuggingFace Transformers config or children classes.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name_or_path: str,
|
||||
normalize_embeddings: bool = True,
|
||||
use_fp16: bool = True,
|
||||
use_bf16: bool = False,
|
||||
query_instruction_for_retrieval: Optional[str] = None,
|
||||
query_instruction_format: str = "{}{}", # specify the format of query_instruction_for_retrieval
|
||||
devices: Optional[Union[str, int, List[str], List[int]]] = None,
|
||||
# inference
|
||||
batch_size: int = 256,
|
||||
query_max_length: int = 512,
|
||||
passage_max_length: int = 512,
|
||||
convert_to_numpy: bool = True,
|
||||
truncate_dim: Optional[int] = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
self.model_name_or_path = model_name_or_path
|
||||
self.normalize_embeddings = normalize_embeddings
|
||||
self.use_fp16 = use_fp16
|
||||
self.use_bf16 = use_bf16
|
||||
self.query_instruction_for_retrieval = query_instruction_for_retrieval
|
||||
self.query_instruction_format = query_instruction_format
|
||||
self.target_devices = self.get_target_devices(devices)
|
||||
|
||||
self.batch_size = batch_size
|
||||
self.query_max_length = query_max_length
|
||||
self.passage_max_length = passage_max_length
|
||||
self.convert_to_numpy = convert_to_numpy
|
||||
self.truncate_dim = truncate_dim
|
||||
|
||||
for k in kwargs:
|
||||
setattr(self, k, kwargs[k])
|
||||
|
||||
self.kwargs = kwargs
|
||||
|
||||
# tokenizer and model are initialized in the child class
|
||||
self.tokenizer = None
|
||||
self.model = None
|
||||
self.pool = None
|
||||
|
||||
def get_model_torch_dtype(self) -> torch.dtype:
|
||||
if self.use_bf16:
|
||||
return torch.bfloat16
|
||||
if self.use_fp16:
|
||||
return torch.float16
|
||||
return torch.float32
|
||||
|
||||
def stop_self_pool(self):
|
||||
if self.pool is not None:
|
||||
self.stop_multi_process_pool(self.pool)
|
||||
self.pool = None
|
||||
try:
|
||||
self.model.to('cpu')
|
||||
torch.cuda.empty_cache()
|
||||
except:
|
||||
pass
|
||||
if gc is not None and callable(gc.collect):
|
||||
gc.collect()
|
||||
|
||||
@staticmethod
|
||||
def get_target_devices(devices: Union[str, int, List[str], List[int]]) -> List[str]:
|
||||
"""
|
||||
|
||||
Args:
|
||||
devices (Union[str, int, List[str], List[int]]): specified devices, can be `str`, `int`, list of `str`, or list of `int`.
|
||||
|
||||
Raises:
|
||||
ValueError: Devices should be a string or an integer or a list of strings or a list of integers.
|
||||
|
||||
Returns:
|
||||
List[str]: A list of target devices in format.
|
||||
"""
|
||||
if devices is None:
|
||||
if torch.cuda.is_available():
|
||||
return [f"cuda:{i}" for i in range(torch.cuda.device_count())]
|
||||
elif is_torch_npu_available():
|
||||
return [f"npu:{i}" for i in range(torch.npu.device_count())]
|
||||
elif hasattr(torch, "musa") and torch.musa.is_available():
|
||||
return [f"musa:{i}" for i in range(torch.musa.device_count())]
|
||||
elif torch.backends.mps.is_available():
|
||||
try:
|
||||
return [f"mps:{i}" for i in range(torch.mps.device_count())]
|
||||
except:
|
||||
return ["mps"]
|
||||
else:
|
||||
return ["cpu"]
|
||||
elif isinstance(devices, str):
|
||||
return [devices]
|
||||
elif isinstance(devices, int):
|
||||
if hasattr(torch, "musa") and torch.musa.is_available():
|
||||
return [f"musa:{devices}"]
|
||||
else:
|
||||
return [f"cuda:{devices}"]
|
||||
elif isinstance(devices, list):
|
||||
if isinstance(devices[0], str):
|
||||
return devices
|
||||
elif isinstance(devices[0], int):
|
||||
if hasattr(torch, "musa") and torch.musa.is_available():
|
||||
return [f"musa:{device}" for device in devices]
|
||||
else:
|
||||
return [f"cuda:{device}" for device in devices]
|
||||
else:
|
||||
raise ValueError("devices should be a string or an integer or a list of strings or a list of integers.")
|
||||
else:
|
||||
raise ValueError("devices should be a string or an integer or a list of strings or a list of integers.")
|
||||
|
||||
@staticmethod
|
||||
def get_detailed_instruct(instruction_format: str, instruction: str, sentence: str):
|
||||
"""Combine the instruction and sentence along with the instruction format.
|
||||
|
||||
Args:
|
||||
instruction_format (str): Format for instruction.
|
||||
instruction (str): The text of instruction.
|
||||
sentence (str): The sentence to concatenate with.
|
||||
|
||||
Returns:
|
||||
str: The complete sentence with instruction
|
||||
"""
|
||||
if "\\n" in instruction_format:
|
||||
instruction_format = instruction_format.replace("\\n", "\n")
|
||||
return instruction_format.format(instruction, sentence)
|
||||
|
||||
def encode_queries(
|
||||
self,
|
||||
queries: Union[List[str], str],
|
||||
batch_size: Optional[int] = None,
|
||||
max_length: Optional[int] = None,
|
||||
convert_to_numpy: Optional[bool] = None,
|
||||
**kwargs: Any
|
||||
):
|
||||
"""encode the queries using the instruction if provided.
|
||||
|
||||
Args:
|
||||
queries (Union[List[str], str]): Input queries to encode.
|
||||
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
|
||||
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
|
||||
convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will
|
||||
be a Torch Tensor. Defaults to :data:`None`.
|
||||
|
||||
Returns:
|
||||
Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
|
||||
"""
|
||||
if batch_size is None: batch_size = self.batch_size
|
||||
if max_length is None: max_length = self.query_max_length
|
||||
if convert_to_numpy is None: convert_to_numpy = self.convert_to_numpy
|
||||
|
||||
return self.encode(
|
||||
queries,
|
||||
batch_size=batch_size,
|
||||
max_length=max_length,
|
||||
convert_to_numpy=convert_to_numpy,
|
||||
instruction=self.query_instruction_for_retrieval,
|
||||
instruction_format=self.query_instruction_format,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
def encode_corpus(
|
||||
self,
|
||||
corpus: Union[List[str], str],
|
||||
batch_size: Optional[int] = None,
|
||||
max_length: Optional[int] = None,
|
||||
convert_to_numpy: Optional[bool] = None,
|
||||
**kwargs: Any
|
||||
):
|
||||
"""encode the corpus using the instruction if provided.
|
||||
|
||||
Args:
|
||||
corpus (Union[List[str], str]): Input corpus to encode.
|
||||
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
|
||||
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
|
||||
convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will
|
||||
be a Torch Tensor. Defaults to :data:`None`.
|
||||
|
||||
Returns:
|
||||
Union[torch.Tensor, np.ndarray]: Return the embedding vectors in a numpy array or tensor.
|
||||
"""
|
||||
passage_instruction_for_retrieval = self.kwargs.get("passage_instruction_for_retrieval", None)
|
||||
passage_instruction_format = self.kwargs.get("passage_instruction_format", "{}{}")
|
||||
|
||||
if batch_size is None: batch_size = self.batch_size
|
||||
if max_length is None: max_length = self.passage_max_length
|
||||
if convert_to_numpy is None: convert_to_numpy = self.convert_to_numpy
|
||||
|
||||
return self.encode(
|
||||
corpus,
|
||||
batch_size=batch_size,
|
||||
max_length=max_length,
|
||||
convert_to_numpy=convert_to_numpy,
|
||||
instruction=passage_instruction_for_retrieval,
|
||||
instruction_format=passage_instruction_format,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
def encode(
|
||||
self,
|
||||
sentences: Union[List[str], str],
|
||||
batch_size: Optional[int] = None,
|
||||
max_length: Optional[int] = None,
|
||||
convert_to_numpy: Optional[bool] = None,
|
||||
instruction: Optional[str] = None,
|
||||
instruction_format: Optional[str] = None,
|
||||
**kwargs: Any
|
||||
):
|
||||
"""encode the input sentences with the embedding model.
|
||||
|
||||
Args:
|
||||
sentences (Union[List[str], str]): Input sentences to encode.
|
||||
batch_size (Optional[int], optional): Number of sentences for each iter. Defaults to :data:`None`.
|
||||
max_length (Optional[int], optional): Maximum length of tokens. Defaults to :data:`None`.
|
||||
convert_to_numpy (Optional[bool], optional): If True, the output embedding will be a Numpy array. Otherwise, it will
|
||||
be a Torch Tensor. Defaults to :data:`None`.
|
||||
instruction (Optional[str], optional): The text of instruction. Defaults to :data:`None`.
|
||||
instruction_format (Optional[str], optional): Format for instruction. Defaults to :data:`None`.
|
||||
|
||||
Returns:
|
||||
Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
|
||||
"""
|
||||
if batch_size is None: batch_size = self.batch_size
|
||||
if max_length is None: max_length = self.passage_max_length
|
||||
if convert_to_numpy is None: convert_to_numpy = self.convert_to_numpy
|
||||
|
||||
if instruction is not None:
|
||||
if isinstance(sentences, str):
|
||||
sentences = self.get_detailed_instruct(instruction_format, instruction, sentences)
|
||||
else:
|
||||
sentences = [self.get_detailed_instruct(instruction_format, instruction, sentence) for sentence in
|
||||
sentences]
|
||||
|
||||
if isinstance(sentences, str) or len(self.target_devices) == 1:
|
||||
return self.encode_single_device(
|
||||
sentences,
|
||||
batch_size=batch_size,
|
||||
max_length=max_length,
|
||||
convert_to_numpy=convert_to_numpy,
|
||||
device=self.target_devices[0],
|
||||
**kwargs
|
||||
)
|
||||
|
||||
if self.pool is None:
|
||||
self.pool = self.start_multi_process_pool(AbsEmbedder._encode_multi_process_worker)
|
||||
embeddings = self.encode_multi_process(
|
||||
sentences,
|
||||
self.pool,
|
||||
batch_size=batch_size,
|
||||
max_length=max_length,
|
||||
convert_to_numpy=convert_to_numpy,
|
||||
**kwargs
|
||||
)
|
||||
return embeddings
|
||||
|
||||
def __del__(self):
|
||||
self.stop_self_pool()
|
||||
|
||||
@abstractmethod
|
||||
def encode_single_device(
|
||||
self,
|
||||
sentences: Union[List[str], str],
|
||||
batch_size: int = 256,
|
||||
max_length: int = 512,
|
||||
convert_to_numpy: bool = True,
|
||||
device: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
This method should encode sentences and return embeddings on a single device.
|
||||
"""
|
||||
pass
|
||||
|
||||
# adapted from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L807
|
||||
def start_multi_process_pool(
|
||||
self,
|
||||
process_target_func: Any,
|
||||
) -> Dict[Literal["input", "output", "processes"], Any]:
|
||||
"""
|
||||
Starts a multi-process pool to process the encoding with several independent processes
|
||||
via :meth:`SentenceTransformer.encode_multi_process <sentence_transformers.SentenceTransformer.encode_multi_process>`.
|
||||
|
||||
This method is recommended if you want to encode on multiple GPUs or CPUs. It is advised
|
||||
to start only one process per GPU. This method works together with encode_multi_process
|
||||
and stop_multi_process_pool.
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: A dictionary with the target processes, an input queue, and an output queue.
|
||||
"""
|
||||
if self.model is None:
|
||||
raise ValueError("Model is not initialized.")
|
||||
|
||||
logger.info("Start multi-process pool on devices: {}".format(", ".join(map(str, self.target_devices))))
|
||||
|
||||
self.model.to("cpu")
|
||||
self.model.share_memory()
|
||||
ctx = mp.get_context("spawn")
|
||||
input_queue = ctx.Queue()
|
||||
output_queue = ctx.Queue()
|
||||
processes = []
|
||||
|
||||
for device_id in tqdm(self.target_devices, desc='initial target device'):
|
||||
p = ctx.Process(
|
||||
target=process_target_func,
|
||||
args=(device_id, self, input_queue, output_queue),
|
||||
daemon=True,
|
||||
)
|
||||
p.start()
|
||||
processes.append(p)
|
||||
|
||||
return {"input": input_queue, "output": output_queue, "processes": processes}
|
||||
|
||||
# adapted from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L976
|
||||
@staticmethod
|
||||
def _encode_multi_process_worker(
|
||||
target_device: str, model: 'AbsEmbedder', input_queue: Queue, results_queue: Queue
|
||||
) -> None:
|
||||
"""
|
||||
Internal working process to encode sentences in multi-process setup
|
||||
"""
|
||||
while True:
|
||||
try:
|
||||
chunk_id, sentences, kwargs = (
|
||||
input_queue.get()
|
||||
)
|
||||
embeddings = model.encode_single_device(
|
||||
sentences,
|
||||
device=target_device,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
results_queue.put([chunk_id, embeddings])
|
||||
except queue.Empty:
|
||||
break
|
||||
|
||||
# copied from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L857
|
||||
@staticmethod
|
||||
def stop_multi_process_pool(pool: Dict[Literal["input", "output", "processes"], Any]) -> None:
|
||||
"""
|
||||
Stops all processes started with start_multi_process_pool.
|
||||
|
||||
Args:
|
||||
pool (Dict[str, object]): A dictionary containing the input queue, output queue, and process list.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
for p in pool["processes"]:
|
||||
p.terminate()
|
||||
|
||||
for p in pool["processes"]:
|
||||
p.join()
|
||||
p.close()
|
||||
|
||||
pool["input"].close()
|
||||
pool["output"].close()
|
||||
pool = None
|
||||
|
||||
# adapted from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L877
|
||||
def encode_multi_process(
|
||||
self,
|
||||
sentences: List[str],
|
||||
pool: Dict[Literal["input", "output", "processes"], Any],
|
||||
**kwargs
|
||||
):
|
||||
chunk_size = math.ceil(len(sentences) / len(pool["processes"]))
|
||||
|
||||
input_queue = pool["input"]
|
||||
last_chunk_id = 0
|
||||
chunk = []
|
||||
|
||||
for sentence in sentences:
|
||||
chunk.append(sentence)
|
||||
if len(chunk) >= chunk_size:
|
||||
input_queue.put(
|
||||
[last_chunk_id, chunk, kwargs]
|
||||
)
|
||||
last_chunk_id += 1
|
||||
chunk = []
|
||||
|
||||
if len(chunk) > 0:
|
||||
input_queue.put([last_chunk_id, chunk, kwargs])
|
||||
last_chunk_id += 1
|
||||
|
||||
output_queue = pool["output"]
|
||||
results_list = sorted(
|
||||
[output_queue.get() for _ in trange(last_chunk_id, desc="Chunks")],
|
||||
key=lambda x: x[0],
|
||||
)
|
||||
embeddings = self._concatenate_results_from_multi_process([result[1] for result in results_list])
|
||||
return embeddings
|
||||
|
||||
def _concatenate_results_from_multi_process(self, results_list: List[Union[torch.Tensor, np.ndarray, Any]]):
|
||||
"""concatenate and return the results from all the processes
|
||||
|
||||
Args:
|
||||
results_list (List[Union[torch.Tensor, np.ndarray, Any]]): A list of results from all the processes.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Unsupported type for results_list
|
||||
|
||||
Returns:
|
||||
Union[torch.Tensor, np.ndarray]: return the embedding vectors in a numpy array or tensor.
|
||||
"""
|
||||
if isinstance(results_list[0], torch.Tensor):
|
||||
# move all tensors to the same device
|
||||
results_list = [res.to(self.target_devices[0]) for res in results_list]
|
||||
return torch.cat(results_list, dim=0)
|
||||
elif isinstance(results_list[0], np.ndarray):
|
||||
return np.concatenate(results_list, axis=0)
|
||||
else:
|
||||
raise NotImplementedError("Unsupported type for results_list")
|
||||
|
||||
def _convert_to_numpy(self, embeddings: torch.Tensor, device: Optional[str] = None) -> np.ndarray:
|
||||
"""Convert tensor embeddings to numpy with bf16-safe handling.
|
||||
|
||||
NumPy does not support bfloat16, so we upcast to float32 only when
|
||||
bf16 inference is enabled on non-CPU devices.
|
||||
|
||||
Args:
|
||||
embeddings (torch.Tensor): Embedding tensor.
|
||||
device (Optional[str], optional): Inference device string. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Embeddings in numpy format.
|
||||
"""
|
||||
if device != "cpu" and self.use_bf16 and embeddings.dtype == torch.bfloat16:
|
||||
embeddings = embeddings.float()
|
||||
return embeddings.cpu().numpy()
|
||||
|
||||
def _truncate_embeddings(self, embeddings: torch.Tensor) -> torch.Tensor:
|
||||
"""Truncate the embedding vectors to the specified dimension.
|
||||
|
||||
This is useful for Matryoshka Representation Learning models, where
|
||||
embeddings can be truncated to a smaller dimension without significant
|
||||
loss of quality.
|
||||
|
||||
Args:
|
||||
embeddings (torch.Tensor): The embedding tensor to truncate.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The truncated embedding tensor. If :attr:`truncate_dim` is None,
|
||||
the original embeddings are returned unchanged.
|
||||
"""
|
||||
if self.truncate_dim is not None:
|
||||
embeddings = embeddings[..., :self.truncate_dim]
|
||||
return embeddings
|
||||
@@ -0,0 +1,360 @@
|
||||
import logging
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Union, List, Tuple, Dict, Literal, Optional
|
||||
|
||||
import multiprocessing as mp
|
||||
from multiprocessing import Queue
|
||||
|
||||
import math
|
||||
import gc
|
||||
import torch
|
||||
import numpy as np
|
||||
from tqdm import tqdm, trange
|
||||
from transformers import is_torch_npu_available
|
||||
|
||||
try:
|
||||
import torch_musa
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsReranker(ABC):
|
||||
"""
|
||||
Base class for Reranker.
|
||||
Extend this class and implement :meth:`compute_score_single_gpu` for custom rerankers.
|
||||
|
||||
Args:
|
||||
model_name_or_path (str): If it's a path to a local model, it loads the model from the path. Otherwise tries to download and
|
||||
load a model from HuggingFace Hub with the name.
|
||||
use_fp16 (bool, optional): If true, use half-precision floating-point to speed up computation with a slight performance
|
||||
degradation. Defaults to :data:`False`.
|
||||
query_instruction_for_rerank: (Optional[str], optional): Query instruction for reranking, which will be used with
|
||||
with :attr:`query_instruction_format`. Defaults to :data:`None`.
|
||||
query_instruction_format: (str, optional): The template for :attr:`query_instruction_for_rerank`. Defaults to :data:`"{}{}"`.
|
||||
passage_instruction_for_rerank (Optional[str], optional): Passage instruction for reranking. Defaults to :data:`None`.
|
||||
passage_instruction_format (str, optional): Passage instruction format when using :attr:`passage_instruction_for_rerank`.
|
||||
Defaults to :data:`"{}{}"`.
|
||||
devices (Optional[Union[str, int, List[str], List[int]]], optional): Devices to use for model inference. Defaults to :data:`None`.
|
||||
batch_size (int, optional): Batch size for inference. Defaults to :data:`128`.
|
||||
query_max_length (int, optional): Maximum length for query. Defaults to :data:`None`.
|
||||
max_length (int, optional): Maximum length. Defaults to :data:`512`.
|
||||
normalize (bool, optional): If true, normalize the result. Defaults to :data:`False`.
|
||||
kwargs (Dict[Any], optional): Additional parameters for HuggingFace Transformers config or children classes.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model_name_or_path: str,
|
||||
use_fp16: bool = False,
|
||||
query_instruction_for_rerank: Optional[str] = None,
|
||||
query_instruction_format: str = "{}{}", # specify the format of query_instruction_for_rerank
|
||||
passage_instruction_for_rerank: Optional[str] = None,
|
||||
passage_instruction_format: str = "{}{}", # specify the format of passage_instruction_for_rerank
|
||||
devices: Optional[Union[str, int, List[str], List[int]]] = None,
|
||||
# inference
|
||||
batch_size: int = 128,
|
||||
query_max_length: Optional[int] = None,
|
||||
max_length: int = 512,
|
||||
normalize: bool = False,
|
||||
**kwargs: Any,
|
||||
):
|
||||
self.model_name_or_path = model_name_or_path
|
||||
self.use_fp16 = use_fp16
|
||||
self.query_instruction_for_rerank = query_instruction_for_rerank
|
||||
self.query_instruction_format = query_instruction_format
|
||||
self.passage_instruction_for_rerank = passage_instruction_for_rerank
|
||||
self.passage_instruction_format = passage_instruction_format
|
||||
self.target_devices = self.get_target_devices(devices)
|
||||
|
||||
self.batch_size = batch_size
|
||||
self.query_max_length = query_max_length
|
||||
self.max_length = max_length
|
||||
self.normalize = normalize
|
||||
|
||||
for k in kwargs:
|
||||
setattr(self, k, kwargs[k])
|
||||
|
||||
self.kwargs = kwargs
|
||||
|
||||
# tokenizer and model are initialized in the child class
|
||||
self.model = None
|
||||
self.tokenizer = None
|
||||
self.pool = None
|
||||
|
||||
def stop_self_pool(self):
|
||||
if self.pool is not None:
|
||||
self.stop_multi_process_pool(self.pool)
|
||||
self.pool = None
|
||||
try:
|
||||
self.model.to('cpu')
|
||||
torch.cuda.empty_cache()
|
||||
except:
|
||||
pass
|
||||
if gc is not None and callable(gc.collect):
|
||||
gc.collect()
|
||||
|
||||
@staticmethod
|
||||
def get_target_devices(devices: Union[str, int, List[str], List[int]]) -> List[str]:
|
||||
"""
|
||||
|
||||
Args:
|
||||
devices (Union[str, int, List[str], List[int]]): Specified devices, can be `str`, `int`, list of `str`, or list of `int`.
|
||||
|
||||
Raises:
|
||||
ValueError: Devices should be a string or an integer or a list of strings or a list of integers.
|
||||
|
||||
Returns:
|
||||
List[str]: A list of target devices in format
|
||||
"""
|
||||
if devices is None:
|
||||
if torch.cuda.is_available():
|
||||
return [f"cuda:{i}" for i in range(torch.cuda.device_count())]
|
||||
elif is_torch_npu_available():
|
||||
return [f"npu:{i}" for i in range(torch.npu.device_count())]
|
||||
elif hasattr(torch, "musa") and torch.musa.is_available():
|
||||
return [f"musa:{i}" for i in range(torch.musa.device_count())]
|
||||
elif torch.backends.mps.is_available():
|
||||
return ["mps"]
|
||||
else:
|
||||
return ["cpu"]
|
||||
elif isinstance(devices, str):
|
||||
return [devices]
|
||||
elif isinstance(devices, int):
|
||||
if hasattr(torch, "musa") and torch.musa.is_available():
|
||||
return [f"musa:{devices}"]
|
||||
else:
|
||||
return [f"cuda:{devices}"]
|
||||
elif isinstance(devices, list):
|
||||
if isinstance(devices[0], str):
|
||||
return devices
|
||||
elif isinstance(devices[0], int):
|
||||
if hasattr(torch, "musa") and torch.musa.is_available():
|
||||
return [f"musa:{device}" for device in devices]
|
||||
else:
|
||||
return [f"cuda:{device}" for device in devices]
|
||||
else:
|
||||
raise ValueError("devices should be a string or an integer or a list of strings or a list of integers.")
|
||||
else:
|
||||
raise ValueError("devices should be a string or an integer or a list of strings or a list of integers.")
|
||||
|
||||
def get_detailed_instruct(self, instruction_format: str, instruction: str, sentence: str):
|
||||
"""Combine the instruction and sentence along with the instruction format.
|
||||
|
||||
Args:
|
||||
instruction_format (str): Format for instruction.
|
||||
instruction (str): The text of instruction.
|
||||
sentence (str): The sentence to concatenate with.
|
||||
|
||||
Returns:
|
||||
str: The complete sentence with instruction
|
||||
"""
|
||||
if "\\n" in instruction_format:
|
||||
instruction_format = instruction_format.replace("\\n", "\n")
|
||||
return instruction_format.format(instruction, sentence)
|
||||
|
||||
def get_detailed_inputs(self, sentence_pairs: Union[str, List[str]]):
|
||||
"""get detailed instruct for all the inputs
|
||||
|
||||
Args:
|
||||
sentence_pairs (Union[str, List[str]]): Input sentence pairs
|
||||
|
||||
Returns:
|
||||
list[list[str]]: The complete sentence pairs with instruction
|
||||
"""
|
||||
if isinstance(sentence_pairs, str):
|
||||
sentence_pairs = [sentence_pairs]
|
||||
|
||||
if self.query_instruction_for_rerank is not None:
|
||||
if self.passage_instruction_for_rerank is None:
|
||||
return [
|
||||
[
|
||||
self.get_detailed_instruct(self.query_instruction_format, self.query_instruction_for_rerank, sentence_pair[0]),
|
||||
sentence_pair[1]
|
||||
] for sentence_pair in sentence_pairs
|
||||
]
|
||||
else:
|
||||
return [
|
||||
[
|
||||
self.get_detailed_instruct(self.query_instruction_format, self.query_instruction_for_rerank, sentence_pair[0]),
|
||||
self.get_detailed_instruct(self.passage_instruction_format, self.passage_instruction_for_rerank, sentence_pair[1])
|
||||
] for sentence_pair in sentence_pairs
|
||||
]
|
||||
else:
|
||||
if self.passage_instruction_for_rerank is None:
|
||||
return [
|
||||
[
|
||||
sentence_pair[0],
|
||||
sentence_pair[1]
|
||||
] for sentence_pair in sentence_pairs
|
||||
]
|
||||
else:
|
||||
return [
|
||||
[
|
||||
sentence_pair[0],
|
||||
self.get_detailed_instruct(self.passage_instruction_format, self.passage_instruction_for_rerank, sentence_pair[1])
|
||||
] for sentence_pair in sentence_pairs
|
||||
]
|
||||
|
||||
def compute_score(
|
||||
self,
|
||||
sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
|
||||
**kwargs
|
||||
):
|
||||
"""Compute score for each sentence pair
|
||||
|
||||
Args:
|
||||
sentence_pairs (Union[List[Tuple[str, str]], Tuple[str, str]]): Input sentence pairs to compute.
|
||||
|
||||
Returns:
|
||||
numpy.ndarray: scores of all the sentence pairs.
|
||||
"""
|
||||
if isinstance(sentence_pairs[0], str):
|
||||
sentence_pairs = [sentence_pairs]
|
||||
sentence_pairs = self.get_detailed_inputs(sentence_pairs)
|
||||
|
||||
if isinstance(sentence_pairs, str) or len(self.target_devices) == 1:
|
||||
return self.compute_score_single_gpu(
|
||||
sentence_pairs,
|
||||
device=self.target_devices[0],
|
||||
**kwargs
|
||||
)
|
||||
|
||||
if self.pool is None:
|
||||
self.pool = self.start_multi_process_pool()
|
||||
scores = self.encode_multi_process(sentence_pairs,
|
||||
self.pool,
|
||||
**kwargs)
|
||||
return scores
|
||||
|
||||
def __del__(self):
|
||||
self.stop_self_pool()
|
||||
|
||||
@abstractmethod
|
||||
def compute_score_single_gpu(
|
||||
self,
|
||||
sentence_pairs: Union[List[Tuple[str, str]], Tuple[str, str]],
|
||||
batch_size: int = 256,
|
||||
query_max_length: Optional[int] = None,
|
||||
max_length: int = 512,
|
||||
normalize: bool = False,
|
||||
device: Optional[str] = None,
|
||||
**kwargs: Any,
|
||||
):
|
||||
"""
|
||||
This method should compute the scores of sentence_pair and return scores.
|
||||
"""
|
||||
pass
|
||||
|
||||
# copied from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L857
|
||||
def start_multi_process_pool(self) -> Dict[Literal["input", "output", "processes"], Any]:
|
||||
"""
|
||||
Starts a multi-process pool to process the encoding with several independent processes
|
||||
via :meth:`SentenceTransformer.encode_multi_process <sentence_transformers.SentenceTransformer.encode_multi_process>`.
|
||||
|
||||
This method is recommended if you want to encode on multiple GPUs or CPUs. It is advised
|
||||
to start only one process per GPU. This method works together with encode_multi_process
|
||||
and stop_multi_process_pool.
|
||||
|
||||
Returns:
|
||||
Dict[str, Any]: A dictionary with the target processes, an input queue, and an output queue.
|
||||
"""
|
||||
logger.info("Start multi-process pool on devices: {}".format(", ".join(map(str, self.target_devices))))
|
||||
|
||||
self.model.to("cpu")
|
||||
self.model.share_memory()
|
||||
ctx = mp.get_context("spawn")
|
||||
input_queue = ctx.Queue()
|
||||
output_queue = ctx.Queue()
|
||||
processes = []
|
||||
|
||||
for device_id in tqdm(self.target_devices, desc='initial target device'):
|
||||
p = ctx.Process(
|
||||
target=AbsReranker._encode_multi_process_worker,
|
||||
args=(device_id, self, input_queue, output_queue),
|
||||
daemon=True,
|
||||
)
|
||||
p.start()
|
||||
processes.append(p)
|
||||
|
||||
return {"input": input_queue, "output": output_queue, "processes": processes}
|
||||
|
||||
# copied from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L857
|
||||
def encode_multi_process(
|
||||
self,
|
||||
sentence_pairs: List,
|
||||
pool: Dict[Literal["input", "output", "processes"], Any],
|
||||
**kwargs
|
||||
) -> np.ndarray:
|
||||
chunk_size = math.ceil(len(sentence_pairs) / len(pool["processes"]))
|
||||
|
||||
input_queue = pool["input"]
|
||||
last_chunk_id = 0
|
||||
chunk = []
|
||||
|
||||
for sentence_pair in sentence_pairs:
|
||||
chunk.append(sentence_pair)
|
||||
if len(chunk) >= chunk_size:
|
||||
input_queue.put(
|
||||
[last_chunk_id, chunk, kwargs]
|
||||
)
|
||||
last_chunk_id += 1
|
||||
chunk = []
|
||||
|
||||
if len(chunk) > 0:
|
||||
input_queue.put([last_chunk_id, chunk, kwargs])
|
||||
last_chunk_id += 1
|
||||
|
||||
output_queue = pool["output"]
|
||||
results_list = sorted(
|
||||
[output_queue.get() for _ in trange(last_chunk_id, desc="Chunks")],
|
||||
key=lambda x: x[0],
|
||||
)
|
||||
scores = np.concatenate([result[1] for result in results_list])
|
||||
return scores
|
||||
|
||||
# copied from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L857
|
||||
@staticmethod
|
||||
def _encode_multi_process_worker(
|
||||
target_device: str, model: 'AbsReranker', input_queue: Queue, results_queue: Queue
|
||||
) -> None:
|
||||
"""
|
||||
Internal working process to encode sentences in multi-process setup
|
||||
"""
|
||||
while True:
|
||||
try:
|
||||
chunk_id, sentences, kwargs = (
|
||||
input_queue.get()
|
||||
)
|
||||
embeddings = model.compute_score_single_gpu(
|
||||
sentences,
|
||||
device=target_device,
|
||||
**kwargs
|
||||
)
|
||||
|
||||
results_queue.put([chunk_id, embeddings])
|
||||
except:
|
||||
break
|
||||
|
||||
# copied from https://github.com/UKPLab/sentence-transformers/blob/1802076d4eae42ff0a5629e1b04e75785d4e193b/sentence_transformers/SentenceTransformer.py#L857
|
||||
@staticmethod
|
||||
def stop_multi_process_pool(pool: Dict[Literal["input", "output", "processes"], Any]) -> None:
|
||||
"""
|
||||
Stops all processes started with start_multi_process_pool.
|
||||
|
||||
Args:
|
||||
pool (Dict[str, object]): A dictionary containing the input queue, output queue, and process list.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
for p in pool["processes"]:
|
||||
p.terminate()
|
||||
|
||||
for p in pool["processes"]:
|
||||
p.join()
|
||||
p.close()
|
||||
|
||||
pool["input"].close()
|
||||
pool["output"].close()
|
||||
@@ -0,0 +1,7 @@
|
||||
from .AbsEmbedder import AbsEmbedder
|
||||
from .AbsReranker import AbsReranker
|
||||
|
||||
__all__ = [
|
||||
'AbsEmbedder',
|
||||
'AbsReranker'
|
||||
]
|
||||
Reference in New Issue
Block a user