File size: 1,678 Bytes
f54cf3c |
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 |
from argparse import ArgumentParser, Namespace
from typing import List, Optional
from model_api import SionicEmbeddingModel
from mteb import MTEB
RETRIEVAL_TASKS: List[str] = [
'ArguAna',
'ClimateFEVER',
'DBPedia',
'FEVER',
'FiQA2018',
'HotpotQA',
'MSMARCO',
'NFCorpus',
'NQ',
'QuoraRetrieval',
'SCIDOCS',
'SciFact',
'Touche2020',
'TRECCOVID',
]
def get_arguments() -> Namespace:
parser = ArgumentParser()
parser.add_argument('--url', type=str, default='https://api.sionic.ai/v2/embedding', help='api server url')
parser.add_argument('--instruction', type=str, default='query: ', help='query instruction')
parser.add_argument('--batch_size', type=int, default=128)
parser.add_argument('--dimension', type=int, default=3072)
parser.add_argument('--output_dir', type=str, default='./result/v2')
return parser.parse_args()
if __name__ == '__main__':
args = get_arguments()
model = SionicEmbeddingModel(url=args.url, instruction=args.instruction, batch_size=args.batch_size, dimension=args.dimension)
task_names: List[str] = [t.description['name'] for t in MTEB(task_types=None, task_langs=['en']).tasks]
for task in task_names:
if task in ['MSMARCOv2']:
continue
instruction: Optional[str] = args.instruction if ('CQADupstack' in task) or (task in RETRIEVAL_TASKS) else None
model.instruction = instruction
evaluation = MTEB(
tasks=[task],
task_langs=['en'],
eval_splits=['test' if task not in ['MSMARCO'] else 'dev'],
)
evaluation.run(model, output_folder=args.output_dir)
|