vlm/engine/runner/loops.py

# Copyright (c) OpenMMLab. All rights reserved.
import copy

from mmengine.runner import ValLoop as MMENGINE_ValLoop
from mmengine.dist import broadcast_object_list, is_main_process, get_world_size, get_rank, barrier, collect_results
import math
import torch
from mmengine.model import is_model_wrapper
from types import MethodType
from xtuner.utils import (DEFAULT_IMAGE_TOKEN, IMAGE_TOKEN_INDEX,
                          PROMPT_TEMPLATE)
from xtuner.tools.utils import get_stop_criteria, is_cn_string
from transformers import GenerationConfig

TORCH_DTYPE_MAP = dict(
    fp16=torch.float16, bf16=torch.bfloat16, fp32=torch.float32, auto='auto')

class TestLoop(MMENGINE_ValLoop):
    def __init__(self, runner, dataloader, evaluator=None, torch_dtype='fp16', select_metric='first') -> None:
        # must be concatset
        super(MMENGINE_ValLoop, self).__init__(runner, dataloader)
        self._runner = runner
        self.torch_dtype = torch_dtype
        if torch_dtype is not None:
            self.torch_dtype = TORCH_DTYPE_MAP[torch_dtype]
        self.select_metric = select_metric

    def run(self) -> dict:
        """Launch Test."""
        self.runner.logger.info('==================== Start test loop ===================')
        self.runner.call_hook('before_test')
        self.runner.call_hook('before_test_epoch')

        if is_model_wrapper(self.runner.model):
            model = self.runner.model.module
        else:
            model = self.runner.model

        model.gradient_checkpointing_disable()
        model.eval()
        model.cuda()

        rank = get_rank()
        metrics = []
        # Ensure that eta and log are displayed correctly.
        current_run_total_ids = 0
        for _, dataset in enumerate(self.dataloader.dataset.datasets):
            if not hasattr(model, 'preparing_for_generation'):
                model.preparing_for_generation = MethodType(default_preparing_for_generation, model)
                print("Warning, the model do not have the preparing_for_generation() function, using the default!!!")
            model.preparing_for_generation(dataset.metainfo)

            # split per rank
            results = []
            n_samples = len(dataset)
            per_rank_samples = math.ceil(n_samples / get_world_size())
            per_rank_ids = range(per_rank_samples * rank,
                                 min(n_samples, per_rank_samples * (rank + 1)))
            for idx in per_rank_ids:
                data_batch = dataset[idx]
                self.run_iter(current_run_total_ids, data_batch, results, model)
                current_run_total_ids += 1

            barrier()
            self.runner.logger.info('==================== Start collect results ===================')
            results = collect_results(results, len(dataset))
            self.runner.logger.info('========= Starting the evaluation of a data ===========')
            if is_main_process():
                metric = dataset.evaluate(results, self.runner.work_dir)
                objects = [metric]
            else:
                objects = [None]
            broadcast_object_list(objects)
            metric = objects[0]
            metrics.append(metric)

        # select metrics
        if self.select_metric == 'first':
            metrics = metrics[0]
        else:
            raise NotImplementedError

        self.runner.logger.info('================ Ending test loop ================')
        self.runner.call_hook('after_test_epoch', metrics=metrics)
        self.runner.call_hook('after_test')
        return metrics

    @torch.no_grad()
    def run_iter(self, idx, data_batch, results, model):
        assert 'text_prompts' in data_batch and 'pixel_values' in data_batch and 'img_id' in data_batch
        prediction = {'img_id': data_batch['img_id']}

        self.runner.call_hook(
            'before_test_iter', batch_idx=idx, data_batch=data_batch)

        outputs = model.predict_forward(**data_batch)
        prediction.update(outputs)
        results.append(prediction)

        self.runner.call_hook(
            'after_test_iter',
            batch_idx=idx,
            data_batch=data_batch,
            outputs=outputs)

def default_preparing_for_generation(self, metainfo):
    # set stop criteria and generation configs for model

    assert hasattr(self, 'tokenizer'), "The Model does not have the tokenizer!!!"

    self.bot_name = 'BOT'
    template = PROMPT_TEMPLATE['internlm2_chat']
    self.template = template
    stop_words = []
    stop_words += template.get('STOP_WORDS', [])
    stop_criteria = get_stop_criteria(
        tokenizer=self.tokenizer, stop_words=stop_words)
    self.stop_criteria = stop_criteria

    default_generation_kwargs = dict(
        max_new_tokens=2048,
        do_sample=False,
        eos_token_id=self.tokenizer.eos_token_id,
        pad_token_id=(
            self.tokenizer.pad_token_id
            if self.tokenizer.pad_token_id is not None
            else self.tokenizer.eos_token_id
        ),
    )
    default_generation_kwargs.update(metainfo.get('generation_kwargs', {}))
    self.gen_config = GenerationConfig(**default_generation_kwargs)
    return


class AnnoLoop(MMENGINE_ValLoop):
    def __init__(self, runner, dataloader, evaluator=None, torch_dtype='fp16', select_metric='first') -> None:
        # must be concatset
        super(MMENGINE_ValLoop, self).__init__(runner, dataloader)
        self._runner = runner
        self.torch_dtype = torch_dtype
        if torch_dtype is not None:
            self.torch_dtype = TORCH_DTYPE_MAP[torch_dtype]
        self.select_metric = select_metric

    def run(self) -> dict:
        """Launch Test."""
        self.runner.logger.info('==================== Start test loop ===================')
        self.runner.call_hook('before_test')
        self.runner.call_hook('before_test_epoch')

        if is_model_wrapper(self.runner.model):
            model = self.runner.model.module
        else:
            model = self.runner.model

        model.eval()

        rank = get_rank()
        metrics = []
        # Ensure that eta and log are displayed correctly.
        current_run_total_ids = 0
        for _, dataset in enumerate(self.dataloader.dataset.datasets):

            # split per rank
            results = []
            n_samples = len(dataset)
            per_rank_samples = math.ceil(n_samples / get_world_size())
            per_rank_ids = range(per_rank_samples * rank,
                                 min(n_samples, per_rank_samples * (rank + 1)))
            for idx in per_rank_ids:
                data_batch = dataset[idx]
                self.run_iter(current_run_total_ids, data_batch, results, model)
                current_run_total_ids += 1
            if hasattr(model, 'save_step'):
                model.save_step(last=True)

            barrier()
            self.runner.logger.info('==================== Start collect results ===================')
            results = collect_results(results, len(dataset))
            self.runner.logger.info('========= Starting the evaluation of a data ===========')
            if is_main_process():
                metric = dataset.evaluate(results, self.runner.work_dir)
                objects = [metric]
            else:
                objects = [None]
            broadcast_object_list(objects)
            metric = objects[0]
            metrics.append(metric)

        # select metrics
        if self.select_metric == 'first':
            metrics = metrics[0]
        else:
            raise NotImplementedError

        self.runner.logger.info('================ Ending test loop ================')
        self.runner.call_hook('after_test_epoch', metrics=metrics)
        self.runner.call_hook('after_test')
        return metrics

    @torch.no_grad()
    def run_iter(self, idx, data_batch, results, model):
        prediction = {}

        self.runner.call_hook(
            'before_test_iter', batch_idx=idx, data_batch=data_batch)

        outputs = model.predict_forward(**data_batch)
        prediction.update(outputs)
        results.append(prediction)

        self.runner.call_hook(
            'after_test_iter',
            batch_idx=idx,
            data_batch=data_batch,
            outputs=outputs)