|
import os
|
|
|
|
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
|
|
os.environ["OPENBLAS_NUM_THREADS"] = "32"
|
|
import mteb
|
|
import torch
|
|
import numpy as np
|
|
from mteb.encoder_interface import PromptType
|
|
from sentence_transformers import SentenceTransformer
|
|
|
|
TASK_NAME2TYPE = {
|
|
'ArguAna': 'Retrieval', 'ArXivHierarchicalClusteringP2P': 'Clustering',
|
|
'ArXivHierarchicalClusteringS2S': 'Clustering', 'AskUbuntuDupQuestions': 'Reranking',
|
|
'BIOSSES': 'STS', 'Banking77Classification': 'Classification',
|
|
'BiorxivClusteringP2P.v2': 'Clustering', 'CQADupstackGamingRetrieval': 'Retrieval',
|
|
'CQADupstackUnixRetrieval': 'Retrieval', 'ClimateFEVERHardNegatives': 'Retrieval',
|
|
'FEVERHardNegatives': 'Retrieval', 'FiQA2018': 'Retrieval', 'HotpotQAHardNegatives': 'Retrieval',
|
|
'ImdbClassification': 'Classification', 'MTOPDomainClassification': 'Classification',
|
|
'MassiveIntentClassification': 'Classification', 'MassiveScenarioClassification': 'Classification',
|
|
'MedrxivClusteringP2P.v2': 'Clustering', 'MedrxivClusteringS2S.v2': 'Clustering',
|
|
'MindSmallReranking': 'Reranking', 'SCIDOCS': 'Retrieval', 'SICK-R': 'STS', 'STS12': 'STS',
|
|
'STS13': 'STS', 'STS14': 'STS', 'STS15': 'STS', 'STSBenchmark': 'STS',
|
|
'SprintDuplicateQuestions': 'PairClassification', 'StackExchangeClustering.v2': 'Clustering',
|
|
'StackExchangeClusteringP2P.v2': 'Clustering', 'TRECCOVID': 'Retrieval',
|
|
'Touche2020Retrieval.v3': 'Retrieval', 'ToxicConversationsClassification': 'Classification',
|
|
'TweetSentimentExtractionClassification': 'Classification',
|
|
'TwentyNewsgroupsClustering.v2': 'Clustering', 'TwitterSemEval2015': 'PairClassification',
|
|
'TwitterURLCorpus': 'PairClassification', 'SummEvalSummarization.v2': 'Summarization',
|
|
'AmazonCounterfactualClassification': 'Classification', 'STS17': 'STS', 'STS22.v2': 'STS'
|
|
}
|
|
|
|
RETRIEVE_Q_PROMPT = "<|START_INSTRUCTION|>Answer the question<|END_INSTRUCTION|>"
|
|
RETRIEVE_P_PROMPT = "<|START_INSTRUCTION|>Candidate document<|END_INSTRUCTION|>"
|
|
STS_PROMPT = "<|START_INSTRUCTION|>Generate semantically similar text<|END_INSTRUCTION|>"
|
|
|
|
TASK_NAME2PROMPT = {
|
|
|
|
"Banking77Classification": "<|START_INSTRUCTION|>Classify text into intents<|END_INSTRUCTION|>",
|
|
"ImdbClassification": "<|START_INSTRUCTION|>Classify text into sentiment<|END_INSTRUCTION|>",
|
|
"MTOPDomainClassification": "<|START_INSTRUCTION|>Classify text into intent domain<|END_INSTRUCTION|>",
|
|
"MassiveIntentClassification": "<|START_INSTRUCTION|>Classify text into user intents<|END_INSTRUCTION|>",
|
|
"MassiveScenarioClassification": "<|START_INSTRUCTION|>Classify text into user scenarios<|END_INSTRUCTION|>",
|
|
"ToxicConversationsClassification": "<|START_INSTRUCTION|>Classify text into toxic or not toxic<|END_INSTRUCTION|>",
|
|
"TweetSentimentExtractionClassification": "<|START_INSTRUCTION|>Classify text into positive, negative, or neutral sentiment<|END_INSTRUCTION|>",
|
|
"AmazonCounterfactualClassification": "<|START_INSTRUCTION|>Classify text into counterfactual or not-counterfactual<|END_INSTRUCTION|>",
|
|
|
|
|
|
"ArXivHierarchicalClusteringP2P": "<|START_INSTRUCTION|>Output main and secondary category of Arxiv papers based on the titles and abstracts<|END_INSTRUCTION|>",
|
|
"ArXivHierarchicalClusteringS2S": "<|START_INSTRUCTION|>Output main and secondary category of Arxiv papers based on the titles<|END_INSTRUCTION|>",
|
|
"BiorxivClusteringP2P.v2": "<|START_INSTRUCTION|>Output main category of Biorxiv papers based on the titles and abstracts<|END_INSTRUCTION|>",
|
|
"MedrxivClusteringP2P.v2": "<|START_INSTRUCTION|>Output main category of Medrxiv papers based on the titles and abstracts<|END_INSTRUCTION|>",
|
|
"MedrxivClusteringS2S.v2": "<|START_INSTRUCTION|>Output main category of Medrxiv papers based on the titles<|END_INSTRUCTION|>",
|
|
"StackExchangeClustering.v2": "<|START_INSTRUCTION|>Output topic or theme of StackExchange posts based on the titles<|END_INSTRUCTION|>",
|
|
"StackExchangeClusteringP2P.v2": "<|START_INSTRUCTION|>Output topic or theme of StackExchange posts based on the given paragraphs<|END_INSTRUCTION|>",
|
|
"TwentyNewsgroupsClustering.v2": "<|START_INSTRUCTION|>Output topic or theme of news articles<|END_INSTRUCTION|>",
|
|
}
|
|
|
|
|
|
class DeweyWrapper:
|
|
def __init__(self, model_dir, max_seq_length: int = 1536, batch_size: int = 8):
|
|
self.model = SentenceTransformer(
|
|
model_dir,
|
|
trust_remote_code=True,
|
|
model_kwargs={
|
|
"torch_dtype": torch.bfloat16,
|
|
"attn_implementation": "flash_attention_2"
|
|
},
|
|
config_kwargs={"single_vector_type": "cls_add_mean"}
|
|
).cuda().bfloat16().eval()
|
|
self.model.max_seq_length = max_seq_length
|
|
self.pool = self.model.start_multi_process_pool()
|
|
self.batch_size = batch_size
|
|
|
|
def encode(
|
|
self,
|
|
sentences: list[str],
|
|
task_name: str,
|
|
prompt_type: PromptType | None = None,
|
|
**kwargs,
|
|
) -> np.ndarray:
|
|
task_type = TASK_NAME2TYPE[task_name]
|
|
if task_type == "Retrieval":
|
|
if prompt_type.value == "query":
|
|
prompt = RETRIEVE_Q_PROMPT
|
|
else:
|
|
prompt = RETRIEVE_P_PROMPT
|
|
elif task_type in ["STS", "PairClassification", "Summarization", "Reranking"]:
|
|
prompt = STS_PROMPT
|
|
else:
|
|
prompt = TASK_NAME2PROMPT[task_name]
|
|
vectors = self.model.encode_multi_process(
|
|
sentences=sentences,
|
|
pool=self.pool,
|
|
show_progress_bar=True,
|
|
batch_size=self.batch_size,
|
|
normalize_embeddings=True,
|
|
prompt=prompt,
|
|
precision="float32"
|
|
)
|
|
return vectors
|
|
|
|
|
|
if __name__ == "__main__":
|
|
max_seq_length = 1536
|
|
batch_szie = 8
|
|
model_dir_or_name = "infgrad/dewey_en_beta"
|
|
output_folder = f"./mteb_eng_results/dewey_en_beta"
|
|
model = DeweyWrapper(model_dir_or_name, max_seq_length=max_seq_length, batch_size=batch_szie)
|
|
|
|
tasks = list(mteb.get_benchmark("MTEB(eng, v2)"))
|
|
evaluation = mteb.MTEB(tasks=tasks)
|
|
evaluation.run(model, output_folder=output_folder, verbosity=2, overwrite_results=False)
|
|
|
