Spaces:
Sleeping
Sleeping
File size: 5,322 Bytes
6fc683c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 |
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, always_add_eos=(args.pool_type == 'last'))
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()
|