File size: 1,598 Bytes
50c36f4 |
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 |
# --------------------------------------------------------
# Copyright (c) 2025 NVIDIA
# Licensed under customized NSCLv1 [see LICENSE.md for details]
# --------------------------------------------------------
import mteb
import argparse
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument(
'--model_name_or_path',
type=str,
help='Path to model checkpoint if HF',
default=''
)
parser.add_argument(
'--task',
type=str,
help='Task to evaluate, if None then whole VisualDocumentRetrieval',
default=None,
)
args, extra_args = parser.parse_known_args()
def convert_value(value):
if value.replace('.', '', 1).isdigit(): # Check if it's a number (int or float)
return int(value) if '.' not in value else float(value)
return value # Keep as string if not numeric
# Convert extra_args list to dictionary with proper type conversion
extra_args_dict = {extra_args[i].lstrip('-'): convert_value(extra_args[i + 1])
for i in range(0, len(extra_args), 2)}
return args, extra_args_dict
if __name__ == "__main__":
args, add_args = get_args()
model_meta = mteb.get_model_meta(args.model_name_or_path)
model = model_meta.load_model()
model.mteb_model_meta = model_meta
if args.task is not None:
tasks = [mteb.get_task(args.task)]
else:
tasks = mteb.get_benchmark("VisualDocumentRetrieval")
evaluation = mteb.MTEB(tasks=tasks)
results = evaluation.run(model, corpus_chunk_size=250) |