content: # Copyright (c) 2023 OpenGVLab
  # Copyright (c) 2025 Bytedance Ltd. and/or its affiliates.
  # SPDX-License-Identifier: MIT
  #
  # This file has been modified by ByteDance Ltd. and/or its affiliates. on 2025-05-20.
  #
  # Original file was released under MIT, with the full license text
  # available at https://github.com/OpenGVLab/InternVL/blob/main/LICENSE.
  #
  # This modified file is released under the same license.
  
  import argparse
  import json
  import os
  import random
  
  import torch
  from eval.vlm.utils import load_model_and_tokenizer, build_transform, process_conversation
  from PIL import Image
  from tqdm import tqdm
  
  ds_collections = {
      'mmvet': {
          'root': 'eval/vlm/data/mm-vet/images',
          'question': 'eval/vlm/data/mm-vet/llava-mm-vet.jsonl',
          'metric': None,
          'max_new_tokens': 1000,
          'min_new_tokens': 1,
      }
  }
  
  
  class VQADataset(torch.utils.data.Dataset):
  
      def __init__(self, root, data, prompt):
          self.root = root
          self.data = open(data).readlines()
          self.prompt = prompt
          
      def __len__(self):
          return len(self.data)
  
      def __getitem__(self, idx):
          data = json.loads(self.data[idx].strip())
          image, question, question_id, annotation = data['image'], data[
              'text'], data['question_id'], data.get('answer', None)
  
          image = os.path.join(self.root, image)
          image = Image.open(image).convert('RGB')
          images = [image]
          
          question = question + ' ' + self.prompt
  
          images, conversation = process_conversation(images, question)
  
          return question_id, question, images, conversation, annotation
  
  
  def evaluate_chat_model():
      random.seed(args.seed)
      prompt = ''
  
      for ds_name in args.datasets:
          dataset = VQADataset(
              root=ds_collections[ds_name]['root'],
              data=ds_collections[ds_name]['question'],
              prompt=prompt,
          )
  
          outputs = {}
          for _, (question_id, question, images, conversation, annotations) in tqdm(enumerate(dataset)):
              pred = model.chat(
                  tokenizer, 
                  new_token_ids,
                  image_transform,
                  images=images,
                  prompt=conversation,
                  max_length=ds_collections[ds_name]['max_new_tokens'], # TODO: how to use ds_collections[ds_name]['min_new_tokens']
              )
  
              outputs[f'v1_{question_id}'] = pred
  
          print(f'Evaluating {ds_name} ...')
          results_file = os.path.join(args.out_dir, 'results.json')
          json.dump(outputs, open(results_file, 'w'))
          print('Results saved to {}'.format(results_file))
  
  
  if __name__ == '__main__':
      parser = argparse.ArgumentParser()
      parser.add_argument('--datasets', type=str, default='mmvet')
      parser.add_argument('--batch-size', type=int, default=1)
      parser.add_argument('--num-workers', type=int, default=1)
      parser.add_argument('--out-dir', type=str, default='results')
      parser.add_argument('--seed', type=int, default=0)
      parser.add_argument('--model-path', type=str, default='hf/BAGEL-7B-MoT/')
      args = parser.parse_args()
  
      if not os.path.exists(args.out_dir):
          os.makedirs(args.out_dir, exist_ok=True)
  
      args.datasets = args.datasets.split(',')
      print('datasets:', args.datasets)
      assert args.batch_size == 1, 'Only batch size 1 is supported'
  
      model, tokenizer, new_token_ids = load_model_and_tokenizer(args)
      image_transform = build_transform()
  
      total_params = sum(p.numel() for p in model.parameters()) / 1e9
      print(f'[test] total_params: {total_params}B')
  
      evaluate_chat_model()
  