Upload mmbench.py

modified_xtuner/xtuner/tools/mmbench.py

@@ -0,0 +1,525 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+import argparse
+import json
+import math
+import os
+import os.path as osp
+import re
+import string
+import time
+
+import numpy as np
+import pandas as pd
+import torch
+import tqdm
+from huggingface_hub import snapshot_download
+from mmengine import mkdir_or_exist
+from mmengine.dist import (collect_results, get_dist_info, get_rank, init_dist,
+                           master_only)
+from mmengine.utils.dl_utils import set_multi_processing
+from peft import PeftModel
+from rich.console import Console
+from rich.table import Table
+from torch.utils.data import Dataset
+from transformers import (AutoModel, AutoModelForCausalLM, AutoTokenizer,
+                          BitsAndBytesConfig, SiglipImageProcessor,
+                          SiglipVisionModel, Dinov2Model,
+                          GenerationConfig)
+
+from xtuner.dataset.utils import decode_base64_to_image, expand2square
+from xtuner.model.utils import LoadWoInit, prepare_inputs_labels_for_multimodal
+from xtuner.tools.utils import get_stop_criteria, is_cn_string
+from xtuner.utils import (DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX,
+                          PROMPT_TEMPLATE)
+
+TORCH_DTYPE_MAP = dict(
+    fp16=torch.float16, bf16=torch.bfloat16, fp32=torch.float32, auto='auto')
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='MMBench')
+    parser.add_argument(
+        'model_name_or_path', help='Hugging Face model name or path')
+    parser.add_argument('--data-path', default=None, help='data path')
+    parser.add_argument('--work-dir', help='the dir to save results')
+    parser.add_argument('--llava', default=None, help='llava name or path')
+    parser.add_argument(
+        '--siglip', default=None, help='siglip visual encoder name or path')
+    parser.add_argument(
+        '--visual-select-layer', default=-2, help='visual select layer')
+    parser.add_argument(
+        '--dino', default=None, help='dino visual encoder name or path')
+    parser.add_argument(
+        '--prompt-template',
+        choices=PROMPT_TEMPLATE.keys(),
+        default=None,
+        help='Specify a prompt template')
+    parser.add_argument(
+        '--stop-words', nargs='+', type=str, default=[], help='Stop words')
+    parser.add_argument(
+        '--torch-dtype',
+        default='fp16',
+        choices=TORCH_DTYPE_MAP.keys(),
+        help='Override the default `torch.dtype` and load the model under '
+        'a specific `dtype`.')
+    parser.add_argument(
+        '--bits',
+        type=int,
+        choices=[4, 8, None],
+        default=None,
+        help='LLM bits')
+    parser.add_argument(
+        '--bot-name', type=str, default='BOT', help='Name for Bot')
+    parser.add_argument(
+        '--offload-folder',
+        default=None,
+        help='The folder in which to offload the model weights (or where the '
+        'model weights are already offloaded).')
+    parser.add_argument(
+        '--max-new-tokens',
+        type=int,
+        default=100,
+        help='Maximum number of new tokens allowed in generated text')
+    parser.add_argument(
+        '--seed',
+        type=int,
+        default=0,
+        help='Random seed for reproducible text generation')
+    parser.add_argument(
+        '--launcher',
+        choices=['none', 'pytorch', 'slurm', 'mpi'],
+        default='none',
+        help='job launcher')
+    args = parser.parse_args()
+    return args
+
+
+@master_only
+def master_print(msg):
+    print(msg)
+
+
+class MMBenchDataset(Dataset):
+    ABBRS = {
+        'coarse_perception': 'CP',
+        'finegrained_perception (instance-level)': 'FP-S',
+        'finegrained_perception (cross-instance)': 'FP-C',
+        'logic_reasoning': 'LR',
+        'relation_reasoning': 'RR',
+        'attribute_reasoning': 'AR',
+        'sketch_reasoning': 'Sketch Reasoning',
+        'scenery_building': 'Scenery & Building',
+        'food_clothes': 'Food & Clothes',
+        'historical_figure': 'Historical Figure',
+        'traditional_show': 'Traditional Show',
+        'calligraphy_painting': 'Calligraphy Painting',
+        'cultural_relic': 'Cultural Relic'
+    }
+
+    def __init__(self, data_file):
+        self.data_file = data_file
+        self.df = pd.read_csv(data_file, sep='\t')
+        self.split = 'dev' if 'answer' in self.df.iloc[0].keys() else 'test'
+        self.has_l2_category = 'l2-category' in self.df.columns.to_list()
+
+    def get_image(self, image):
+        while len(image) < 16:
+            image = self.df[self.df['index'] == int(image)]['image'].values
+            assert len(image) == 1
+            image = image[0]
+        image = decode_base64_to_image(image)
+        return image
+
+    def __len__(self):
+        return len(self.df)
+
+    def __getitem__(self, idx):
+        index = self.df.iloc[idx]['index']
+        image = self.df.iloc[idx]['image']
+        image = self.get_image(image)
+        question = self.df.iloc[idx]['question']
+        answer = self.df.iloc[idx]['answer'] if 'answer' in self.df.iloc[
+            0].keys() else None
+        category = self.df.iloc[idx]['category']
+
+        options = {
+            cand: self.load_from_df(idx, cand)
+            for cand in string.ascii_uppercase
+            if self.load_from_df(idx, cand) is not None
+        }
+        options_prompt = ''
+        for key, item in options.items():
+            options_prompt += f'{key}. {item}\n'
+
+        hint = self.load_from_df(idx, 'hint')
+        data = {
+            'img': image,
+            'question': question,
+            'answer': answer,
+            'options': options_prompt,
+            'category': category,
+            'options_dict': options,
+            'index': index,
+            'context': hint,
+        }
+        if self.has_l2_category:
+            data.update({'l2-category': self.df.iloc[idx]['l2-category']})
+        return data
+
+    def load_from_df(self, idx, key):
+        if key in self.df.iloc[idx] and not pd.isna(self.df.iloc[idx][key]):
+            return self.df.iloc[idx][key]
+        else:
+            return None
+
+    @master_only
+    def eval_result(self, result_df, show=True):
+
+        def calc_acc(df, group='category'):
+            assert group in ['overall', 'category', 'l2-category']
+            if group == 'overall':
+                res = {'Average': np.mean(df['hit'])}
+            else:
+                res = {}
+                abilities = list(set(df[group]))
+                abilities.sort()
+                for ab in abilities:
+                    sub_df = df[df[group] == ab]
+                    ab = self.ABBRS[ab] if ab in self.ABBRS else ab
+                    res[ab] = np.mean(sub_df['hit'])
+            return res
+
+        def eval_sub_data(sub_data, answer_map):
+            lt = len(sub_data)
+            for i in range(lt):
+                item = sub_data.iloc[i]
+                match = re.search(r'([A-D]+)', item['prediction'])
+                pred = match.group(1) if match else ''
+                gt = answer_map[item['index']]
+                if gt != pred:
+                    return 0
+            return 1
+
+        def show_result(ret_json):
+            show_dict = ret_json.copy()
+            table = Table(title=f' MMBench ({self.data_file}) ')
+            console = Console()
+            table.add_column('Category', justify='left')
+            table.add_column('Accuracy (%)', justify='right')
+            average = show_dict.pop('Average') * 100
+            table.add_row('Average', f'{average:.1f}')
+            table.add_section()
+            for cat_name, cat_acc in show_dict.items():
+                table.add_row(cat_name, f'{cat_acc * 100:.1f}')
+            with console.capture() as capture:
+                console.print(table, end='')
+            print('\n' + capture.get())
+            print('Note: Please be cautious if you use the results in papers, '
+                  "since we don't use ChatGPT as a helper for choice "
+                  'extraction')
+
+        data = result_df.sort_values(by='index')
+        data['prediction'] = [str(x) for x in data['prediction']]
+        for k in data.keys():
+            data[k.lower() if k not in 'ABCD' else k] = data.pop(k)
+
+        data_main = data[data['index'] < int(1e6)]
+        cate_map = {
+            i: c
+            for i, c in zip(self.df['index'], self.df['category'])
+        }
+        if self.has_l2_category:
+            l2_cate_map = {
+                i: c
+                for i, c in zip(self.df['index'], self.df['l2-category'])
+            }
+        answer_map = {
+            i: c
+            for i, c in zip(self.df['index'], self.df['answer'])
+        }
+
+        lt = len(data_main)
+        hit, tot = 0, 0
+        result = {}
+        for i in range(lt):
+            item_main = data_main.iloc[i]
+            idx = item_main['index']
+            assert idx not in result
+            sub_data = data[data['index'] % int(1e6) == idx]
+            ret = eval_sub_data(sub_data, answer_map)
+            result[idx] = ret
+            hit += ret
+            tot += 1
+
+        indices = data_main['index']
+        data_main = data_main.copy()
+        data_main['hit'] = [result[i] for i in indices]
+        main_idx = data_main['index']
+        data_main['category'] = [cate_map[i] for i in main_idx]
+
+        ret_json = calc_acc(data_main, 'overall')
+
+        if self.has_l2_category:
+            data_main['l2-category'] = [l2_cate_map[i] for i in main_idx]
+            l2 = calc_acc(data_main, 'l2-category')
+            ret_json.update(l2)
+        else:
+            leaf = calc_acc(data_main, 'category')
+            ret_json.update(leaf)
+        if show:
+            show_result(ret_json)
+        return ret_json
+
+
+def main():
+    args = parse_args()
+
+    torch.manual_seed(args.seed)
+
+    if args.launcher != 'none':
+        set_multi_processing(distributed=True)
+        init_dist(args.launcher)
+
+        rank, world_size = get_dist_info()
+        torch.cuda.set_device(rank)
+    else:
+        rank = 0
+        world_size = 1
+
+    # build llm
+    quantization_config = None
+    load_in_8bit = False
+    if args.bits == 4:
+        quantization_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            load_in_8bit=False,
+            llm_int8_threshold=6.0,
+            llm_int8_has_fp16_weight=False,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type='nf4')
+    elif args.bits == 8:
+        load_in_8bit = True
+    model_kwargs = {
+        'quantization_config': quantization_config,
+        'load_in_8bit': load_in_8bit,
+        'device_map': rank if world_size > 1 else 'auto',
+        'offload_folder': args.offload_folder,
+        'trust_remote_code': True,
+        'torch_dtype': TORCH_DTYPE_MAP[args.torch_dtype]
+    }
+
+    # build llm
+    with LoadWoInit():
+        llm = AutoModelForCausalLM.from_pretrained(args.model_name_or_path,
+                                                   **model_kwargs)
+    tokenizer = AutoTokenizer.from_pretrained(
+        args.model_name_or_path,
+        trust_remote_code=True,
+        encode_special_tokens=True)
+    master_print(f'Load LLM from {args.model_name_or_path}')
+
+    llava_path = snapshot_download(
+        repo_id=args.llava) if not osp.isdir(args.llava) else args.llava
+
+    # build visual_encoder
+    if 'visual_encoder' in os.listdir(llava_path):
+        assert args.visual_encoder is None, (
+            "Please don't specify the `--visual-encoder` since passed "
+            '`--llava` contains a visual encoder!')
+        visual_encoder_path = osp.join(llava_path, 'visual_encoder')
+    else:
+        assert args.siglip is not None, (
+            'Please specify the `--siglip`!')
+        assert args.dino is not None, (
+            'Please specify the `--dino`!')
+    with LoadWoInit():
+        siglip = SiglipVisionModel.from_pretrained(
+            args.siglip,
+            torch_dtype=TORCH_DTYPE_MAP[args.torch_dtype])
+        image_processor = SiglipImageProcessor.from_pretrained(
+            args.siglip)
+        master_print(f'Load siglip from {args.siglip}')
+        dino = Dinov2Model.from_pretrained(
+            args.dino,
+            torch_dtype=TORCH_DTYPE_MAP[args.torch_dtype])
+        master_print(f'Load dino from {args.dino}')
+
+    # load adapter
+    if 'llm_adapter' in os.listdir(llava_path):
+        adapter_path = osp.join(llava_path, 'llm_adapter')
+
+        with LoadWoInit():
+            llm = PeftModel.from_pretrained(
+                llm, adapter_path, offload_folder=args.offload_folder)
+
+        master_print(f'Load LLM adapter from {args.llava}')
+
+    if 'visual_encoder_adapter' in os.listdir(llava_path):
+        adapter_path = osp.join(llava_path, 'visual_encoder_adapter')
+        visual_encoder = PeftModel.from_pretrained(
+            visual_encoder, adapter_path, offload_folder=args.offload_folder)
+        master_print(f'Load visual_encoder adapter from {args.llava}')
+
+    # build projector
+    projector_path = osp.join(llava_path, 'projector')
+    with LoadWoInit():
+        projector = AutoModel.from_pretrained(
+            projector_path, torch_dtype=TORCH_DTYPE_MAP[args.torch_dtype])
+    master_print(f'Load projector from {args.llava}')
+
+    projector.cuda()
+    projector.eval()
+
+    siglip.cuda()
+    siglip.eval()
+    dino.cuda()
+    dino.eval()
+
+    llm.eval()
+
+    stop_words = args.stop_words
+    if args.prompt_template:
+        template = PROMPT_TEMPLATE[args.prompt_template]
+        stop_words += template.get('STOP_WORDS', [])
+    stop_criteria = get_stop_criteria(
+        tokenizer=tokenizer, stop_words=stop_words)
+
+    gen_config = GenerationConfig(
+        max_new_tokens=args.max_new_tokens,
+        do_sample=False,
+        eos_token_id=tokenizer.eos_token_id,
+        pad_token_id=tokenizer.pad_token_id
+        if tokenizer.pad_token_id is not None else tokenizer.eos_token_id,
+    )
+
+    # work_dir
+    if args.work_dir is not None:
+        # update configs according to CLI args if args.work_dir is not None
+        save_dir = args.work_dir
+    else:
+        # use config filename as default work_dir
+        save_dir = osp.join('./work_dirs',
+                            osp.splitext(osp.basename(args.data_path))[0])
+    timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime(time.time()))
+    save_dir = osp.join(save_dir, timestamp)
+
+    if rank == 0:
+        mkdir_or_exist(osp.abspath(save_dir))
+        print('=======================================================')
+        print(f'Dataset path: {osp.abspath(args.data_path)}\n'
+              f'Results will be saved to {osp.abspath(save_dir)}')
+        print('=======================================================')
+
+        args_path = osp.join(save_dir, 'args.json')
+        with open(args_path, 'w') as f:
+            json.dump(args.__dict__, f, indent=2)
+
+    results_xlsx_path = osp.join(save_dir, 'mmbench_result.xlsx')
+    results_json_path = osp.join(save_dir, 'mmbench_result.json')
+
+    dataset = MMBenchDataset(args.data_path)
+
+    results = []
+    n_samples = len(dataset)
+    per_rank_samples = math.ceil(n_samples / world_size)
+
+    per_rank_ids = range(per_rank_samples * rank,
+                         min(n_samples, per_rank_samples * (rank + 1)))
+    for i in tqdm.tqdm(per_rank_ids, desc=f'Rank {rank}'):
+        data_sample = dataset[i]
+        if data_sample['context'] is not None:
+            text = data_sample['context'] + '\n' + data_sample[
+                'question'] + '\n' + data_sample['options']
+        else:
+            text = data_sample['question'] + '\n' + data_sample['options']
+
+        text = DEFAULT_IMAGE_TOKEN + '\n' + text
+
+        if is_cn_string(text):
+            text = text + '请直接回答选项字母。'
+        else:
+            text = text + ("Answer with the option's letter from the "
+                           'given choices directly.')
+
+        if args.prompt_template:
+            prompt_text = ''
+            template = PROMPT_TEMPLATE[args.prompt_template]
+            prompt_text += template['INSTRUCTION'].format(
+                input=text, round=1, bot_name=args.bot_name)
+        else:
+            prompt_text = text
+        inputs = prompt_text
+
+        image = data_sample['img'].convert('RGB')
+        image = expand2square(
+            image, tuple(int(x * 255) for x in image_processor.image_mean))
+        image = image_processor.preprocess(
+            image, return_tensors='pt')['pixel_values'][0]
+        image = image.cuda().unsqueeze(0)
+
+        siglip_out = siglip(
+            image, output_hidden_states=True).hidden_states[args.visual_select_layer]
+        dino_out = dino(
+                image, output_hidden_states=True).hidden_states[-1][:, 1:]
+        visual_out = torch.cat((siglip_out, dino_out), dim=-1)
+        pixel_values = projector(visual_out)
+
+        chunk_encode = []
+        for idx, chunk in enumerate(inputs.split(DEFAULT_IMAGE_TOKEN)):
+            if idx == 0:
+                cur_encode = tokenizer.encode(chunk)
+            else:
+                cur_encode = tokenizer.encode(chunk, add_special_tokens=False)
+            chunk_encode.append(cur_encode)
+        assert len(chunk_encode) == 2
+        ids = []
+        for idx, cur_chunk_encode in enumerate(chunk_encode):
+            ids.extend(cur_chunk_encode)
+            if idx != len(chunk_encode) - 1:
+                ids.append(IMAGE_TOKEN_INDEX)
+        ids = torch.tensor(ids).cuda().unsqueeze(0)
+        mm_inputs = prepare_inputs_labels_for_multimodal(
+            llm=llm, input_ids=ids, pixel_values=pixel_values)
+
+        generate_output = llm.generate(
+            **mm_inputs,
+            generation_config=gen_config,
+            streamer=None,
+            bos_token_id=tokenizer.bos_token_id,
+            stopping_criteria=stop_criteria)
+
+        predict = tokenizer.decode(
+            generate_output[0], skip_special_tokens=True).strip()
+        cur_result = {}
+        cur_result['question'] = data_sample.get('question')
+        cur_result.update(data_sample.get('options_dict'))
+        cur_result['prediction'] = predict
+        if data_sample.get('category') is not None:
+            cur_result['category'] = data_sample.get('category')
+        if data_sample.get('l2-category') is not None:
+            cur_result['l2-category'] = data_sample.get('l2-category')
+        cur_result['index'] = data_sample.get('index')
+        cur_result['split'] = data_sample.get('split')
+        cur_result['answer'] = data_sample.get('answer')
+        results.append(cur_result)
+
+    results = collect_results(results, n_samples)
+
+    if get_rank() == 0:
+
+        results_df = pd.DataFrame(results)
+        with pd.ExcelWriter(results_xlsx_path, engine='openpyxl') as writer:
+            results_df.to_excel(writer, index=False)
+
+        if dataset.split == 'dev':
+            results_dict = dataset.eval_result(results_df, show=True)
+            with open(results_json_path, 'w') as f:
+                json.dump(results_dict, f, indent=2)
+        else:
+            print('All done!')
+
+
+if __name__ == '__main__':
+
+    main()