Modify eval_mm for MiniCPM-V 2.6

eval_mm/vlmevalkit/vlmeval/inference_video.py

@@ -0,0 +1,166 @@
+import torch
+import torch.distributed as dist
+from vlmeval.config import supported_VLM
+from vlmeval.utils import track_progress_rich
+from vlmeval.smp import *
+
+FAIL_MSG = 'Failed to obtain answer via API.'
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--data', type=str, nargs='+', required=True)
+    parser.add_argument('--model', type=str, nargs='+', required=True)
+    parser.add_argument('--nproc', type=int, default=4, required=True)
+    parser.add_argument('--verbose', action='store_true')
+    args = parser.parse_args()
+    return args
+
+
+# Only API model is accepted
+def infer_data_api(work_dir, model_name, dataset, nframe=8, pack=False, samples_dict={}, api_nproc=4):
+    rank, world_size = get_rank_and_world_size()
+    assert rank == 0 and world_size == 1
+    dataset_name = dataset.dataset_name
+    model = supported_VLM[model_name]() if isinstance(model_name, str) else model_name
+    assert getattr(model, 'is_api', False)
+
+    indices = list(samples_dict.keys())
+    structs = [dataset.build_prompt(samples_dict[idx], num_frames=nframe,
+                                    video_llm=getattr(model, 'VIDEO_LLM', False)) for idx in indices]
+
+    packstr = 'pack' if pack else 'nopack'
+    out_file = f'{work_dir}/{model_name}_{dataset_name}_{nframe}frame_{packstr}_supp.pkl'
+    res = load(out_file) if osp.exists(out_file) else {}
+
+    structs = [s for i, s in zip(indices, structs) if i not in res]
+    indices = [i for i in indices if i not in res]
+
+    gen_func = model.generate
+    structs = [dict(message=struct, dataset=dataset_name) for struct in structs]
+
+    if len(structs):
+        track_progress_rich(gen_func, structs, nproc=api_nproc, chunksize=api_nproc, save=out_file, keys=indices)
+
+    res = load(out_file)
+    return res
+
+
+def infer_data(model_name, work_dir, dataset, out_file, nframe=8, pack=False, verbose=False, api_nproc=4):
+    res = load(out_file) if osp.exists(out_file) else {}
+    rank, world_size = get_rank_and_world_size()
+    dataset_name = dataset.dataset_name
+
+    sample_indices = list(dataset.videos) if pack else list(dataset.data['index'])
+    samples = list(dataset.videos) if pack else list(range(len(dataset.data)))
+    sample_map = {i: s for i, s in zip(sample_indices, samples)}
+
+    sample_indices_sub = sample_indices[rank::world_size]
+    if np.all([idx in res for idx in sample_indices_sub]):
+        return model_name
+    sample_indices_subrem = [x for x in sample_indices_sub if x not in res]
+
+    model = supported_VLM[model_name]() if isinstance(model_name, str) else model_name
+
+    is_api = getattr(model, 'is_api', False)
+    if is_api:
+        assert world_size == 1
+        supp = infer_data_api(
+            work_dir=work_dir,
+            model_name=model_name,
+            dataset=dataset,
+            nframe=nframe,
+            pack=pack,
+            samples_dict={k: sample_map[k] for k in sample_indices_subrem},
+            api_nproc=api_nproc)
+        for k in sample_indices_subrem:
+            assert k in supp
+        res.update(supp)
+        dump(res, out_file)
+        return model_name
+
+    for i, idx in tqdm(enumerate(sample_indices_subrem)):
+        if idx in res:
+            continue
+        # adapt to model frame sample number first
+        nframe = getattr(model, 'nframe', 0) if getattr(model, 'nframe', 0) > 0 else nframe
+        # when using video-llm, build prompt returns video+question; otherwise, several frames+question
+        struct = dataset.build_prompt(sample_map[idx], num_frames=nframe, video_llm=getattr(model, 'VIDEO_LLM', False))
+        response = model.generate(message=struct, dataset=dataset_name)
+        torch.cuda.empty_cache()
+
+        if verbose:
+            print(response, flush=True)
+
+        res[idx] = response
+        if (i + 1) % 20 == 0:
+            dump(res, out_file)
+
+    res = {k: res[k] for k in sample_indices_sub}
+    dump(res, out_file)
+    return model
+
+
+# A wrapper for infer_data, do the pre & post processing
+def infer_data_job_video(
+        model,
+        work_dir,
+        model_name,
+        dataset,
+        nframe=8,
+        pack=False,
+        verbose=False,
+        subtitle=False,
+        api_nproc=4):
+
+    dataset_name = dataset.dataset_name
+    packstr = 'pack' if pack else 'nopack'
+    rank, world_size = get_rank_and_world_size()
+    result_file = osp.join(work_dir, f'{model_name}_{dataset_name}_{nframe}frame_{packstr}.xlsx')
+    if dataset_name == 'Video-MME':
+        subtitle_str = 'subs' if subtitle else 'nosubs'
+        result_file = result_file.replace('.xlsx', f'_{subtitle_str}.xlsx')
+
+    # Dump Predictions to Prev File if result file exists
+    if osp.exists(result_file):
+        return model_name
+
+    tmpl = osp.join(work_dir, '{}' + f'{world_size}_{dataset_name}_{nframe}frame_{packstr}.pkl')
+    if dataset_name == 'Video-MME':
+        subtitle_str = 'subs' if subtitle else 'nosubs'
+        tmpl = tmpl.replace('.pkl', f'_{subtitle_str}.pkl')
+    out_file = tmpl.format(rank)
+
+    model = infer_data(
+        model,
+        work_dir=work_dir,
+        dataset=dataset,
+        nframe=nframe,
+        pack=pack,
+        out_file=out_file,
+        verbose=verbose,
+        api_nproc=api_nproc)
+
+    if world_size > 1:
+        dist.barrier()
+
+    if rank == 0:
+        data_all = {}
+        for i in range(world_size):
+            data_all.update(load(tmpl.format(i)))
+
+        meta = dataset.data
+        if dataset_name == 'MMBench-Video' and pack:
+            meta, vstats = dataset.load_pack_answers(data_all)
+            print(f'Statitics of Pack Video Inference: {vstats}')
+        else:
+            for x in meta['index']:
+                assert x in data_all
+            meta['prediction'] = [str(data_all[x]) for x in meta['index']]
+            if 'image' in meta:
+                meta.pop('image')
+
+        dump(meta, result_file)
+        for i in range(world_size):
+            os.remove(tmpl.format(i))
+    return model