Spaces:
Running
Running
| import sys | |
| import datetime | |
| import json | |
| import os | |
| import torch | |
| script_dir = os.path.dirname(os.path.realpath(__file__)) | |
| sys.path.append(os.path.join(script_dir, '..')) | |
| from datasets.vqa_dataset import docVQADataset, docVQATESTDataset, textVQADataset | |
| print(torch.__version__) | |
| import numpy as np | |
| from eval_utils.getargs import parse_args | |
| from eval_utils.vqa_evaluate import * | |
| def get_model(args): | |
| if args.model_name == '': | |
| raise Exception('Model name cannot be empty str!') | |
| from models.MiniCPM.minicpmv import MiniCPM_V, MiniCPM_V_2_6 | |
| model_path = args.model_path | |
| ckpt = args.ckpt | |
| if args.model_name == 'minicpmv': | |
| model = MiniCPM_V(model_path=model_path, ckpt=ckpt, device=args.device) | |
| elif args.model_name == 'minicpmv26': | |
| model = MiniCPM_V_2_6(model_path=model_path, ckpt=ckpt, device=args.device) | |
| else: | |
| raise Exception(f"Unexpected Moedel Name {args.model_name}!") | |
| return model | |
| def main(args): | |
| np.random.seed(0) | |
| max_sample_num = None | |
| torch.distributed.init_process_group( | |
| backend='nccl', | |
| world_size=int(os.getenv('WORLD_SIZE', '1')), | |
| rank=int(os.getenv('RANK', '0')), | |
| ) | |
| torch.cuda.set_device(int(os.getenv('LOCAL_RANK', 0))) | |
| print(f'Init Rank-{torch.distributed.get_rank()}') | |
| if torch.distributed.is_initialized(): | |
| args.device = torch.device(f"cuda:{torch.cuda.current_device()}") | |
| model = get_model(args) | |
| result = {} | |
| time = datetime.datetime.now().strftime("%Y%m%d%H%M%S") | |
| if args.eval_textVQA or args.eval_all: | |
| dataset = textVQADataset(args.textVQA_image_dir, args.textVQA_ann_path) | |
| if max_sample_num is not None: | |
| dataset = torch.utils.data.Subset(dataset, range(max_sample_num)) | |
| acc = evaluate_VQA(model, dataset, args.model_name, 'textVQA', time, \ | |
| batch_size=args.batchsize, generate_method=args.generate_method, answer_path=args.answer_path) | |
| result['textVQA'] = acc | |
| if args.eval_docVQA or args.eval_all: | |
| dataset = docVQADataset(args.docVQA_image_dir, args.docVQA_ann_path) | |
| if max_sample_num is not None: | |
| dataset = torch.utils.data.Subset(dataset, range(max_sample_num)) | |
| acc = evaluate_VQA(model, dataset, args.model_name, 'docVQA', time, batch_size=args.batchsize, generate_method=args.generate_method, answer_path=args.answer_path) | |
| result['docVQA'] = acc | |
| if args.eval_docVQATest or args.eval_all: | |
| target_dataset = "docVQATest" | |
| dataset = docVQATESTDataset(args.docVQATest_image_dir, args.docVQATest_ann_path) | |
| if max_sample_num is not None: | |
| dataset = torch.utils.data.Subset(dataset, range(max_sample_num)) | |
| acc = evaluate_VQA(model, dataset, args.model_name, target_dataset, time, batch_size=args.batchsize, generate_method=args.generate_method, answer_path=args.answer_path) | |
| result['docVQATest'] = acc | |
| if torch.distributed.is_initialized(): | |
| torch.distributed.barrier() | |
| if torch.distributed.is_initialized() and torch.distributed.get_rank() != 0: | |
| return None | |
| result_path = os.path.join(os.path.join(args.answer_path, args.model_name), 'result.json') | |
| output_flag = False | |
| for k, v in result.items(): | |
| if v > 0.0: | |
| output_flag = True | |
| break | |
| if output_flag: | |
| with open(result_path, "w") as f: | |
| f.write(json.dumps(result, indent=4)) | |
| if __name__ == "__main__": | |
| args = parse_args() | |
| main(args) |