mirrors/microsoft--unilm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:24:13 +08:00
commit 1037506f2e
6050 changed files with 1731598 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
e5/README.md
+120
View File
@@ -0,0 +1,120 @@
# E5 Text Embeddings
[Multilingual E5 Text Embeddings: A Technical Report](https://arxiv.org/pdf/2402.05672).
Liang Wang, Nan Yang, Xiaolong Huang, Linjun Yang, Rangan Majumder, Furu Wei, arXiv 2024
[Improving Text Embeddings with Large Language Models](https://arxiv.org/pdf/2401.00368.pdf).
Liang Wang, Nan Yang, Xiaolong Huang, Linjun Yang, Rangan Majumder, Furu Wei, arXiv 2024
[Text Embeddings by Weakly-Supervised Contrastive Pre-training](https://arxiv.org/pdf/2212.03533.pdf).
Liang Wang, Nan Yang, Xiaolong Huang, Binxing Jiao, Linjun Yang, Daxin Jiang, Rangan Majumder, Furu Wei, arXiv 2022
## LLM based Models
| | BEIR | # of layers | embedding dimension | Huggingface |
|------------------------|------|:-----------:|:-------------------:|-----------------------------------------------------------------------------------------|
| E5-mistral-7b-instruct | 56.9 | 32 | 4096 | [intfloat/e5-mistral-7b-instruct](https://huggingface.co/intfloat/e5-mistral-7b-instruct)|
## English Pre-trained Models
| | BEIR | # of layers | embedding dimension | Huggingface |
|------------------------|------|:-----------:|:-------------------:|-----------------------------------------------------------------------------------------|
| E5-small-v2 | 49.0 | 12 | 384 | [intfloat/e5-small-v2](https://huggingface.co/intfloat/e5-small-v2) |
| E5-base-v2 | 50.3 | 12 | 768 | [intfloat/e5-base-v2](https://huggingface.co/intfloat/e5-base-v2) |
| E5-large-v2 | 50.6 | 24 | 1024 | [intfloat/e5-large-v2](https://huggingface.co/intfloat/e5-large-v2) |
| | | | | |
| E5-small | 46.0 | 12 | 384 | [intfloat/e5-small](https://huggingface.co/intfloat/e5-small) |
| E5-base | 48.8 | 12 | 768 | [intfloat/e5-base](https://huggingface.co/intfloat/e5-base) |
| E5-large | 50.0 | 24 | 1024 | [intfloat/e5-large](https://huggingface.co/intfloat/e5-large) |
| | | | | |
| E5-small-unsupervised | 40.8 | 12 | 384 | [intfloat/e5-small-unsupervised](https://huggingface.co/intfloat/e5-small-unsupervised) |
| E5-base-unsupervised | 42.9 | 12 | 768 | [intfloat/e5-base-unsupervised](https://huggingface.co/intfloat/e5-base-unsupervised) |
| E5-large-unsupervised | 44.2 | 24 | 1024 | [intfloat/e5-large-unsupervised](https://huggingface.co/intfloat/e5-large-unsupervised) |
The models with `-unsupervised` suffix only pre-trains on unlabeled datasets.
## Multilingual Pre-trained Models
| | BEIR | # of layers | embedding dimension | Huggingface |
|--------------------------------|------|:-----------:|:-------------------:|-----------------------------------------------------------------------------------------------------------|
| multilingual-e5-small | 46.6 | 12 | 384 | [intfloat/multilingual-e5-small](https://huggingface.co/intfloat/multilingual-e5-small) |
| multilingual-e5-base | 48.9 | 12 | 768 | [intfloat/multilingual-e5-base](https://huggingface.co/intfloat/multilingual-e5-base) |
| multilingual-e5-large | 51.4 | 24 | 1024 | [intfloat/multilingual-e5-large](https://huggingface.co/intfloat/multilingual-e5-large) |
| multilingual-e5-large-instruct | 52.5 | 24 | 1024 | [intfloat/multilingual-e5-large-instruct](https://huggingface.co/intfloat/multilingual-e5-large-instruct) |
## Install Python Package Requirements
```shell
pip install -r requirements.txt
```
For `e5-mistral-7b-instruct`, it would require `transformers>=4.34` to load Mistral model.
## Evaluate on the [BEIR Benchmark](https://arxiv.org/abs/2104.08663)
After installing the required python packages,
run the following command on GPU machines:
```shell
bash scripts/eval_mteb_beir.sh intfloat/e5-small-v2
```
By default,
the evaluation script will use all the available GPUs.
Caution: it could take quite a long time (~10 hours) due to corpus encoding.
For `intfloat/e5-mistral-7b-instruct`, it could take even longer (several days).
## Evaluate on the [MTEB Benchmark](https://arxiv.org/abs/2210.07316)
Run the following command:
```shell
bash scripts/eval_mteb_except_retrieval.sh intfloat/e5-small-v2
```
For multilingual models, simply add a `--multilingual` suffix:
```shell
bash scripts/eval_mteb_except_retrieval.sh intfloat/multilingual-e5-base --multilingual
```
## Other Resources
The data for our proposed synthetic task _personalized passkey retrieval_ is available at [https://huggingface.co/datasets/intfloat/personalized_passkey_retrieval](https://huggingface.co/datasets/intfloat/personalized_passkey_retrieval).
## Troubleshooting
If you encounter OOM error, please try to reduce the batch size.
## Citation
If you find our paper or models helpful, please consider cite as follows:
```
@article{wang2024multilingual,
title={Multilingual E5 Text Embeddings: A Technical Report},
author={Wang, Liang and Yang, Nan and Huang, Xiaolong and Yang, Linjun and Majumder, Rangan and Wei, Furu},
journal={arXiv preprint arXiv:2402.05672},
year={2024}
}
@article{wang2023improving,
title={Improving Text Embeddings with Large Language Models},
author={Wang, Liang and Yang, Nan and Huang, Xiaolong and Yang, Linjun and Majumder, Rangan and Wei, Furu},
journal={arXiv preprint arXiv:2401.00368},
year={2023}
}
@article{wang2022text,
title={Text Embeddings by Weakly-Supervised Contrastive Pre-training},
author={Wang, Liang and Yang, Nan and Huang, Xiaolong and Jiao, Binxing and Yang, Linjun and Jiang, Daxin and Majumder, Rangan and Wei, Furu},
journal={arXiv preprint arXiv:2212.03533},
year={2022}
}
```
## License
This project is licensed under the license found in the LICENSE file in the root directory of this source tree.
[Microsoft Open Source Code of Conduct](https://opensource.microsoft.com/codeofconduct)
e5/model_config.py
+34
View File
@@ -0,0 +1,34 @@
MODEL_NAME_TO_POOL_TYPE = {
'e5-small': 'avg',
'e5-base': 'avg',
'e5-large': 'avg',
'e5-small-unsupervised': 'avg',
'e5-base-unsupervised': 'avg',
'e5-large-unsupervised': 'avg',
'e5-small-v2': 'avg',
'e5-base-v2': 'avg',
'e5-large-v2': 'avg',
'multilingual-e5-small': 'avg',
'multilingual-e5-base': 'avg',
'multilingual-e5-large': 'avg',
'multilingual-e5-large-instruct': 'avg',
'e5-mistral-7b-instruct': 'last',
}
MODEL_NAME_TO_PREFIX_TYPE = {
'e5-small': 'query_or_passage',
'e5-base': 'query_or_passage',
'e5-large': 'query_or_passage',
'e5-small-unsupervised': 'query_or_passage',
'e5-base-unsupervised': 'query_or_passage',
'e5-large-unsupervised': 'query_or_passage',
'e5-small-v2': 'query_or_passage',
'e5-base-v2': 'query_or_passage',
'e5-large-v2': 'query_or_passage',
'multilingual-e5-small': 'query_or_passage',
'multilingual-e5-base': 'query_or_passage',
'multilingual-e5-large': 'query_or_passage',
'multilingual-e5-large-instruct': 'instruction',
'e5-mistral-7b-instruct': 'instruction',
}
e5/mteb_beir_eval.py
+121
View File
@@ -0,0 +1,121 @@
import os
import json
import tqdm
import numpy as np
import torch
import argparse
import torch.nn.functional as F
from typing import List, Dict
from transformers import AutoModel, AutoTokenizer
from transformers.modeling_outputs import BaseModelOutput
from mteb import MTEB, AbsTaskRetrieval, DRESModel
from utils import pool, logger, move_to_cuda, get_detailed_instruct, get_task_def_by_task_name_and_type, create_batch_dict
from model_config import MODEL_NAME_TO_POOL_TYPE, MODEL_NAME_TO_PREFIX_TYPE
parser = argparse.ArgumentParser(description='evaluation for BEIR benchmark')
parser.add_argument('--model-name-or-path', default='intfloat/e5-small-v2',
type=str, metavar='N', help='which model to use')
parser.add_argument('--output-dir', default='tmp-outputs/',
type=str, metavar='N', help='output directory')
parser.add_argument('--doc-as-query', action='store_true', help='use query prefix for passages, only used for Quora as it is a symmetric task')
parser.add_argument('--pool-type', default='avg', help='pool type')
parser.add_argument('--prefix-type', default='query_or_passage', help='prefix type')
parser.add_argument('--dry-run', action='store_true', help='whether to run the script in dry run mode')
args = parser.parse_args()
base_name: str = args.model_name_or_path.split('/')[-1]
args.pool_type = MODEL_NAME_TO_POOL_TYPE.get(base_name, args.pool_type)
args.prefix_type = MODEL_NAME_TO_PREFIX_TYPE.get(base_name, args.prefix_type)
logger.info('Args: {}'.format(json.dumps(args.__dict__, ensure_ascii=False, indent=4)))
assert args.pool_type in ['cls', 'avg', 'last', 'weightedavg'], 'pool_type should be cls / avg / last'
assert args.prefix_type in ['query_or_passage', 'instruction'], 'prefix_type should be query_or_passage / instruction'
os.makedirs(args.output_dir, exist_ok=True)
class RetrievalModel(DRESModel):
# Refer to the code of DRESModel for the methods to overwrite
def __init__(self, **kwargs):
self.encoder = AutoModel.from_pretrained(args.model_name_or_path, torch_dtype=torch.float16)
self.tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
self.prompt = None
self.gpu_count = torch.cuda.device_count()
if self.gpu_count > 1:
self.encoder = torch.nn.DataParallel(self.encoder)
self.encoder.cuda()
self.encoder.eval()
def encode_queries(self, queries: List[str], **kwargs) -> np.ndarray:
if args.prefix_type == 'query_or_passage':
input_texts = [f'query: {q}' for q in queries]
else:
input_texts = [self.prompt + q for q in queries]
return self._do_encode(input_texts)
def encode_corpus(self, corpus: List[Dict[str, str]], **kwargs) -> np.ndarray:
if args.doc_as_query:
return self.encode_queries([d['text'] for d in corpus], **kwargs)
input_texts = ['{} {}'.format(doc.get('title', ''), doc['text']).strip() for doc in corpus]
# no need to add prefix for instruct models
if args.prefix_type == 'query_or_passage':
input_texts = ['passage: {}'.format(t) for t in input_texts]
return self._do_encode(input_texts)
@torch.no_grad()
def _do_encode(self, input_texts: List[str]) -> np.ndarray:
encoded_embeds = []
batch_size = 64 * self.gpu_count
for start_idx in tqdm.tqdm(range(0, len(input_texts), batch_size), desc='encoding', mininterval=10):
batch_input_texts: List[str] = input_texts[start_idx: start_idx + batch_size]
batch_dict = create_batch_dict(self.tokenizer, batch_input_texts)
batch_dict = move_to_cuda(batch_dict)
with torch.cuda.amp.autocast():
outputs: BaseModelOutput = self.encoder(**batch_dict)
embeds = pool(outputs.last_hidden_state, batch_dict['attention_mask'], args.pool_type)
embeds = F.normalize(embeds, p=2, dim=-1)
encoded_embeds.append(embeds.cpu().numpy())
return np.concatenate(encoded_embeds, axis=0)
def set_prompt(self, prompt: str):
self.prompt = prompt
def main():
assert AbsTaskRetrieval.is_dres_compatible(RetrievalModel)
model = RetrievalModel()
task_names = [t.description["name"] for t in MTEB(task_types=['Retrieval'], task_langs=['en']).tasks]
task_names = [t for t in task_names if t != 'MSMARCOv2']
logger.info('Tasks: {}'.format(task_names))
for task in task_names:
if args.dry_run and task not in ['SciFact', 'FiQA2018']:
continue
logger.info('Processing task: {}'.format(task))
if args.prefix_type == 'query_or_passage':
args.doc_as_query = task in ['QuoraRetrieval']
else:
task_def: str = get_task_def_by_task_name_and_type(task_name=task, task_type='Retrieval')
prompt: str = get_detailed_instruct(task_def)
model.set_prompt(prompt=prompt)
logger.info('Set prompt: {}'.format(prompt))
evaluation = MTEB(tasks=[task], task_langs=['en'])
evaluation.run(model, eval_splits=["test" if task not in ['MSMARCO'] else 'dev'],
output_folder=args.output_dir)
if __name__ == '__main__':
main()
e5/mteb_except_retrieval_eval.py
+128
View File
@@ -0,0 +1,128 @@
import os
import torch
import torch.nn.functional as F
import tqdm
import json
import numpy as np
import argparse
from transformers import AutoModel, AutoTokenizer
from transformers.modeling_outputs import BaseModelOutput
from typing import List
from mteb import MTEB
from utils import logger, pool, move_to_cuda, get_detailed_instruct, get_task_def_by_task_name_and_type, create_batch_dict
from model_config import MODEL_NAME_TO_POOL_TYPE, MODEL_NAME_TO_PREFIX_TYPE
parser = argparse.ArgumentParser(description='evaluation for MTEB benchmark except its Retrieval category')
parser.add_argument('--task-types', nargs='+', default=[], help='task types to evaluate')
parser.add_argument('--output-dir', default='',
type=str, metavar='N', help='output directory')
parser.add_argument('--model-name-or-path', default='tmp-outputs/',
type=str, metavar='N', help='which model to use')
parser.add_argument('--pool-type', default='avg', help='pool type')
parser.add_argument('--prefix-type', default='query_or_passage', help='prefix type')
parser.add_argument('--multilingual', action='store_true', help='whether to use multilingual model')
parser.add_argument('--dry-run', action='store_true', help='whether to run the script in dry run mode')
args = parser.parse_args()
base_name: str = args.model_name_or_path.split('/')[-1]
args.pool_type = MODEL_NAME_TO_POOL_TYPE.get(base_name, args.pool_type)
args.prefix_type = MODEL_NAME_TO_PREFIX_TYPE.get(base_name, args.prefix_type)
logger.info('Args: {}'.format(json.dumps(args.__dict__, ensure_ascii=False, indent=4)))
assert args.pool_type in ['cls', 'avg', 'last', 'weightedavg'], 'pool_type should be cls / avg / last'
assert args.prefix_type in ['query_or_passage', 'instruction'], 'prefix_type should be query_or_passage / instruction'
os.makedirs(args.output_dir, exist_ok=True)
class DenseEncoder(torch.nn.Module):
def __init__(self, **kwargs):
super().__init__()
self.encoder = AutoModel.from_pretrained(args.model_name_or_path, torch_dtype=torch.float16)
self.tokenizer = AutoTokenizer.from_pretrained(args.model_name_or_path)
self.l2_normalize = True
self.prompt = None
self.gpu_count = torch.cuda.device_count()
self.encoder.eval()
self.encoder.cuda()
if self.gpu_count > 1:
self.encoder = torch.nn.DataParallel(self.encoder)
@torch.no_grad()
def encode(self, sentences, **kwargs) -> np.ndarray:
""" Returns a list of embeddings for the given sentences.
Args:
sentences (`List[str]`): List of sentences to encode
batch_size (`int`): Batch size for the encoding
Returns:
`List[np.ndarray]` or `List[tensor]`: List of embeddings for the given sentences
"""
input_texts: List[str] = [self.prompt + s for s in sentences]
encoded_embeds = []
batch_size = 64 * self.gpu_count
for start_idx in tqdm.tqdm(range(0, len(input_texts), batch_size), desc='encoding', mininterval=10):
batch_input_texts: List[str] = input_texts[start_idx: start_idx + batch_size]
batch_dict = create_batch_dict(self.tokenizer, batch_input_texts)
batch_dict = move_to_cuda(batch_dict)
with torch.cuda.amp.autocast():
outputs: BaseModelOutput = self.encoder(**batch_dict)
embeds = pool(outputs.last_hidden_state, batch_dict['attention_mask'], args.pool_type)
if self.l2_normalize:
embeds = F.normalize(embeds, p=2, dim=-1)
encoded_embeds.append(embeds.cpu().numpy())
return np.concatenate(encoded_embeds, axis=0)
def set_prompt(self, prompt: str):
self.prompt = prompt
def main():
model = DenseEncoder()
args.task_types = [t for t in args.task_types if t.strip()]
evaluation = MTEB(
task_types=args.task_types or None,
task_langs=['en'] if not args.multilingual else None
)
for task_cls in evaluation.tasks:
task_name: str = task_cls.description['name']
task_type: str = task_cls.description['type']
if args.dry_run and task_name not in ['Banking77Classification', 'ImdbClassification', 'STS12']:
continue
if args.prefix_type == 'query_or_passage':
prompt: str = 'query: '
else:
task_def: str = get_task_def_by_task_name_and_type(task_name=task_name, task_type=task_type)
prompt: str = get_detailed_instruct(task_def)
model.set_prompt(prompt=prompt)
logger.info('Set prompt: {}'.format(prompt))
# disable l2 normalize for classification tasks, as it achieves slightly better results
if task_type == 'Classification':
logger.info('Set l2_normalize to False for classification task')
model.l2_normalize = False
else:
model.l2_normalize = True
logger.info('Set l2_normalize to {}'.format(model.l2_normalize))
sub_eval = MTEB(tasks=[task_name], task_langs=['en'] if not args.multilingual else None)
logger.info('Running evaluation for task: {}, type: {}'.format(task_name, task_type))
eval_splits = ["test"] if "test" in task_cls.description["eval_splits"] else task_cls.description["eval_splits"]
sub_eval.run(
model, eval_splits=eval_splits,
output_folder=args.output_dir
)
if __name__ == '__main__':
main()
e5/requirements.txt
+5
View File
@@ -0,0 +1,5 @@
torch>=1.7
transformers>=4.15.0
tqdm
beir
mteb==1.0.2
e5/scripts/eval_mteb_beir.sh
+23
View File
@@ -0,0 +1,23 @@
#!/usr/bin/env bash
set -x
set -e
DIR="$( cd "$( dirname "$0" )" && cd .. && pwd )"
echo "working directory: ${DIR}"
MODEL_NAME_OR_PATH=""
if [[ $# -ge 1 && ! "$1" == "--"* ]]; then
MODEL_NAME_OR_PATH=$1
shift
fi
if [ -z "$OUTPUT_DIR" ]; then
OUTPUT_DIR="tmp-outputs/"
fi
python -u mteb_beir_eval.py \
--model-name-or-path "${MODEL_NAME_OR_PATH}" \
--output-dir "${OUTPUT_DIR}" "$@"
echo "done"
e5/scripts/eval_mteb_except_retrieval.sh
+24
View File
@@ -0,0 +1,24 @@
#!/usr/bin/env bash
set -x
set -e
DIR="$( cd "$( dirname "$0" )" && cd .. && pwd )"
echo "working directory: ${DIR}"
MODEL_NAME_OR_PATH=""
if [[ $# -ge 1 && ! "$1" == "--"* ]]; then
MODEL_NAME_OR_PATH=$1
shift
fi
if [ -z "$OUTPUT_DIR" ]; then
OUTPUT_DIR="tmp-outputs/"
fi
python -u mteb_except_retrieval_eval.py \
--model-name-or-path "${MODEL_NAME_OR_PATH}" \
--task-types "STS" "Summarization" "PairClassification" "Classification" "Reranking" "Clustering" "BitextMining" \
--output-dir "${OUTPUT_DIR}" "$@"
echo "done"
e5/utils.py
+211
View File
@@ -0,0 +1,211 @@
import torch
import logging
from torch import Tensor
from transformers import PreTrainedTokenizerFast, BatchEncoding
from typing import Mapping, Dict, List
def _setup_logger():
log_format = logging.Formatter("[%(asctime)s %(levelname)s] %(message)s")
logger = logging.getLogger()
logger.setLevel(logging.INFO)
console_handler = logging.StreamHandler()
console_handler.setFormatter(log_format)
logger.handlers = [console_handler]
return logger
logger = _setup_logger()
def move_to_cuda(sample):
if len(sample) == 0:
return {}
def _move_to_cuda(maybe_tensor):
if torch.is_tensor(maybe_tensor):
return maybe_tensor.cuda(non_blocking=True)
elif isinstance(maybe_tensor, dict):
return {key: _move_to_cuda(value) for key, value in maybe_tensor.items()}
elif isinstance(maybe_tensor, list):
return [_move_to_cuda(x) for x in maybe_tensor]
elif isinstance(maybe_tensor, tuple):
return tuple([_move_to_cuda(x) for x in maybe_tensor])
elif isinstance(maybe_tensor, Mapping):
return type(maybe_tensor)({k: _move_to_cuda(v) for k, v in maybe_tensor.items()})
else:
return maybe_tensor
return _move_to_cuda(sample)
def pool(last_hidden_states: Tensor,
attention_mask: Tensor,
pool_type: str) -> Tensor:
last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)
if pool_type == "avg":
emb = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
elif pool_type == "weightedavg": # position-weighted mean pooling from SGPT (https://arxiv.org/abs/2202.08904)
attention_mask *= attention_mask.cumsum(dim=1) # [0,1,1,1,0,0] -> [0,1,2,3,0,0]
s = torch.sum(last_hidden * attention_mask.unsqueeze(-1).float(), dim=1)
d = attention_mask.sum(dim=1, keepdim=True).float()
emb = s / d
elif pool_type == "cls":
emb = last_hidden[:, 0]
elif pool_type == "last":
left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0])
if left_padding:
emb = last_hidden[:, -1]
else:
sequence_lengths = attention_mask.sum(dim=1) - 1
batch_size = last_hidden.shape[0]
emb = last_hidden[torch.arange(batch_size, device=last_hidden.device), sequence_lengths]
else:
raise ValueError(f"pool_type {pool_type} not supported")
return emb
def create_batch_dict(tokenizer: PreTrainedTokenizerFast, input_texts: List[str], max_length: int = 512) -> BatchEncoding:
return tokenizer(
input_texts,
max_length=max_length,
padding=True,
pad_to_multiple_of=8,
return_token_type_ids=False,
truncation=True,
return_tensors='pt'
)
def get_task_def_by_task_name_and_type(task_name: str, task_type: str) -> str:
if task_type in ['STS']:
return "Retrieve semantically similar text."
if task_type in ['Summarization']:
return "Given a news summary, retrieve other semantically similar summaries"
if task_type in ['BitextMining']:
return "Retrieve parallel sentences."
if task_type in ['Classification']:
task_name_to_instruct: Dict[str, str] = {
'AmazonCounterfactualClassification': 'Classify a given Amazon customer review text as either counterfactual or not-counterfactual',
'AmazonPolarityClassification': 'Classify Amazon reviews into positive or negative sentiment',
'AmazonReviewsClassification': 'Classify the given Amazon review into its appropriate rating category',
'Banking77Classification': 'Given a online banking query, find the corresponding intents',
'EmotionClassification': 'Classify the emotion expressed in the given Twitter message into one of the six emotions: anger, fear, joy, love, sadness, and surprise',
'ImdbClassification': 'Classify the sentiment expressed in the given movie review text from the IMDB dataset',
'MassiveIntentClassification': 'Given a user utterance as query, find the user intents',
'MassiveScenarioClassification': 'Given a user utterance as query, find the user scenarios',
'MTOPDomainClassification': 'Classify the intent domain of the given utterance in task-oriented conversation',
'MTOPIntentClassification': 'Classify the intent of the given utterance in task-oriented conversation',
'ToxicConversationsClassification': 'Classify the given comments as either toxic or not toxic',
'TweetSentimentExtractionClassification': 'Classify the sentiment of a given tweet as either positive, negative, or neutral',
# C-MTEB eval instructions
'TNews': 'Classify the fine-grained category of the given news title',
'IFlyTek': 'Given an App description text, find the appropriate fine-grained category',
'MultilingualSentiment': 'Classify sentiment of the customer review into positive, neutral, or negative',
'JDReview': 'Classify the customer review for iPhone on e-commerce platform into positive or negative',
'OnlineShopping': 'Classify the customer review for online shopping into positive or negative',
'Waimai': 'Classify the customer review from a food takeaway platform into positive or negative',
}
return task_name_to_instruct[task_name]
if task_type in ['Clustering']:
task_name_to_instruct: Dict[str, str] = {
'ArxivClusteringP2P': 'Identify the main and secondary category of Arxiv papers based on the titles and abstracts',
'ArxivClusteringS2S': 'Identify the main and secondary category of Arxiv papers based on the titles',
'BiorxivClusteringP2P': 'Identify the main category of Biorxiv papers based on the titles and abstracts',
'BiorxivClusteringS2S': 'Identify the main category of Biorxiv papers based on the titles',
'MedrxivClusteringP2P': 'Identify the main category of Medrxiv papers based on the titles and abstracts',
'MedrxivClusteringS2S': 'Identify the main category of Medrxiv papers based on the titles',
'RedditClustering': 'Identify the topic or theme of Reddit posts based on the titles',
'RedditClusteringP2P': 'Identify the topic or theme of Reddit posts based on the titles and posts',
'StackExchangeClustering': 'Identify the topic or theme of StackExchange posts based on the titles',
'StackExchangeClusteringP2P': 'Identify the topic or theme of StackExchange posts based on the given paragraphs',
'TwentyNewsgroupsClustering': 'Identify the topic or theme of the given news articles',
# C-MTEB eval instructions
'CLSClusteringS2S': 'Identify the main category of scholar papers based on the titles',
'CLSClusteringP2P': 'Identify the main category of scholar papers based on the titles and abstracts',
'ThuNewsClusteringS2S': 'Identify the topic or theme of the given news articles based on the titles',
'ThuNewsClusteringP2P': 'Identify the topic or theme of the given news articles based on the titles and contents',
}
return task_name_to_instruct[task_name]
if task_type in ['Reranking', 'PairClassification']:
task_name_to_instruct: Dict[str, str] = {
'AskUbuntuDupQuestions': 'Retrieve duplicate questions from AskUbuntu forum',
'MindSmallReranking': 'Retrieve relevant news articles based on user browsing history',
'SciDocsRR': 'Given a title of a scientific paper, retrieve the titles of other relevant papers',
'StackOverflowDupQuestions': 'Retrieve duplicate questions from StackOverflow forum',
'SprintDuplicateQuestions': 'Retrieve duplicate questions from Sprint forum',
'TwitterSemEval2015': 'Retrieve tweets that are semantically similar to the given tweet',
'TwitterURLCorpus': 'Retrieve tweets that are semantically similar to the given tweet',
# C-MTEB eval instructions
'T2Reranking': 'Given a Chinese search query, retrieve web passages that answer the question',
'MMarcoReranking': 'Given a Chinese search query, retrieve web passages that answer the question',
'CMedQAv1': 'Given a Chinese community medical question, retrieve replies that best answer the question',
'CMedQAv2': 'Given a Chinese community medical question, retrieve replies that best answer the question',
'Ocnli': 'Retrieve semantically similar text.',
'Cmnli': 'Retrieve semantically similar text.',
}
return task_name_to_instruct[task_name]
if task_type in ['Retrieval']:
if task_name.lower().startswith('cqadupstack'):
return 'Given a question, retrieve detailed question descriptions from Stackexchange that are duplicates to the given question'
task_name_to_instruct: Dict[str, str] = {
'ArguAna': 'Given a claim, find documents that refute the claim',
'ClimateFEVER': 'Given a claim about climate change, retrieve documents that support or refute the claim',
'DBPedia': 'Given a query, retrieve relevant entity descriptions from DBPedia',
'FEVER': 'Given a claim, retrieve documents that support or refute the claim',
'FiQA2018': 'Given a financial question, retrieve user replies that best answer the question',
'HotpotQA': 'Given a multi-hop question, retrieve documents that can help answer the question',
'MSMARCO': 'Given a web search query, retrieve relevant passages that answer the query',
'NFCorpus': 'Given a question, retrieve relevant documents that best answer the question',
'NQ': 'Given a question, retrieve Wikipedia passages that answer the question',
'QuoraRetrieval': 'Given a question, retrieve questions that are semantically equivalent to the given question',
'SCIDOCS': 'Given a scientific paper title, retrieve paper abstracts that are cited by the given paper',
'SciFact': 'Given a scientific claim, retrieve documents that support or refute the claim',
'Touche2020': 'Given a question, retrieve detailed and persuasive arguments that answer the question',
'TRECCOVID': 'Given a query on COVID-19, retrieve documents that answer the query',
# C-MTEB eval instructions
'T2Retrieval': 'Given a Chinese search query, retrieve web passages that answer the question',
'MMarcoRetrieval': 'Given a web search query, retrieve relevant passages that answer the query',
'DuRetrieval': 'Given a Chinese search query, retrieve web passages that answer the question',
'CovidRetrieval': 'Given a question on COVID-19, retrieve news articles that answer the question',
'CmedqaRetrieval': 'Given a Chinese community medical question, retrieve replies that best answer the question',
'EcomRetrieval': 'Given a user query from an e-commerce website, retrieve description sentences of relevant products',
'MedicalRetrieval': 'Given a medical question, retrieve user replies that best answer the question',
'VideoRetrieval': 'Given a video search query, retrieve the titles of relevant videos',
}
# add lower case keys to match some beir names
task_name_to_instruct.update({k.lower(): v for k, v in task_name_to_instruct.items()})
# other cases where lower case match still doesn't work
task_name_to_instruct['trec-covid'] = task_name_to_instruct['TRECCOVID']
task_name_to_instruct['climate-fever'] = task_name_to_instruct['ClimateFEVER']
task_name_to_instruct['dbpedia-entity'] = task_name_to_instruct['DBPedia']
task_name_to_instruct['webis-touche2020'] = task_name_to_instruct['Touche2020']
task_name_to_instruct['fiqa'] = task_name_to_instruct['FiQA2018']
task_name_to_instruct['quora'] = task_name_to_instruct['QuoraRetrieval']
# for miracl evaluation
task_name_to_instruct['miracl'] = 'Given a question, retrieve Wikipedia passages that answer the question'
return task_name_to_instruct[task_name]
raise ValueError(f"No instruction config for task {task_name} with type {task_type}")
def get_detailed_instruct(task_description: str) -> str:
if not task_description:
return ''
return 'Instruct: {}\nQuery: '.format(task_description)