mmpretrain.models.multimodal.ofa.ofa 源代码

# Copyright (c) OpenMMLab. All rights reserved.
import string
from collections import defaultdict
from functools import partial
from typing import Optional, Union

import mmengine
import torch
from mmengine.model import BaseModel

from mmpretrain.datasets import CleanCaption
from mmpretrain.registry import MODELS, TOKENIZER
from mmpretrain.structures import DataSample
from .ofa_modules import OFAEncoderDecoder

class TreeNode():

    def __init__(self):
        self.child = defaultdict(TreeNode)

class Trie:

    def __init__(self, eos):
        self.root = TreeNode()
        self.eos = eos

    def insert(self, word):
        cur = self.root
        for c in word:
            cur = cur.child[c]

    def get_next_layer(self, word):
        cur = self.root
        for c in word:
            cur = cur.child.get(c)
            if cur is None:
                return [self.eos]
        return list(cur.child.keys())

def apply_constraint(
    input_ids: torch.Tensor,
    logits: torch.Tensor,
    decoder_prompts: Optional[list],
    num_beams: int,
    constraint_trie: Trie = None,
    if decoder_prompts is None and constraint_trie is None:
        return logits

    mask = logits.new_zeros(logits[:, -1, :].size(), dtype=torch.bool)
    input_ids = input_ids.view(-1, num_beams, input_ids.shape[-1])
    for batch_id, beam_sent in enumerate(input_ids):
        for beam_id, sent in enumerate(beam_sent):
            if decoder_prompts is None:
                prompt_len = 0
                prompt_len = len(decoder_prompts[batch_id])

            if sent.size(0) - 1 < prompt_len:
                allowed_tokens = [decoder_prompts[batch_id][sent.size(0) - 1]]
                mask[batch_id * num_beams + beam_id, allowed_tokens] = True
            elif constraint_trie is not None:
                answer_tokens = [0] + sent[prompt_len + 1:].tolist()
                allowed_tokens = constraint_trie.get_next_layer(answer_tokens)
                mask[batch_id * num_beams + beam_id, allowed_tokens] = True
                mask[batch_id * num_beams + beam_id, :] = True
    logits[:, -1, :].masked_fill_(~mask, float('-inf'))
    return logits

[文档]@MODELS.register_module() class OFA(BaseModel): """The OFA model for multiple tasks. Args: encoder_cfg (dict): The config of the encoder, accept the keyword arguments of :class:`OFAEncoder`. decoder_cfg (dict): The config of the decoder, accept the keyword arguments of :class:`OFADecoder`. vocab_size (int): The size of the vocabulary. embedding_dim (int): The embedding dimensions of both the encoder and the decoder. tokenizer (dict | PreTrainedTokenizer): The tokenizer to encode the text. task (str): The task name, supported tasks are "caption", "vqa" and "refcoco". prompt (str, optional): The prompt template for the following tasks, If None, use default prompt: - **caption**: ' what does the image describe?' - **refcoco**: ' which region does the text " {} " describe?' Defaults to None ans2label (str | Sequence | None): The answer to label mapping for the vqa task. If a string, it should be a pickle or json file. The sequence constrains the output answers. Defaults to None, which means no constraint. generation_cfg (dict): The extra generation config, accept the keyword arguments of :class:`~transformers.GenerationConfig`. Defaults to an empty dict. data_preprocessor (dict, optional): The config for preprocessing input data. If None or no specified type, it will use "MultiModalDataPreprocessor" as type. See :class: `MultiModalDataPreprocessor` for more details. Defaults to None. init_cfg (dict, optional): The initialization config. Defaults to None. """ support_tasks = {'caption', 'vqa', 'refcoco'} def __init__( self, encoder_cfg, decoder_cfg, vocab_size, embedding_dim, tokenizer, task, prompt=None, ans2label: Union[dict, str, None] = None, generation_cfg=dict(), data_preprocessor: Optional[dict] = None, init_cfg=None, ): if data_preprocessor is None: data_preprocessor = {} if isinstance(data_preprocessor, dict): data_preprocessor.setdefault('type', 'MultiModalDataPreprocessor') data_preprocessor = super().__init__( init_cfg=init_cfg, data_preprocessor=data_preprocessor) if isinstance(tokenizer, dict): self.tokenizer = else: self.tokenizer = tokenizer if task not in self.support_tasks: raise ValueError(f'Unsupported task {task}, please select ' f'the task from {self.support_tasks}.') self.prompt = prompt self.task = task if isinstance(ans2label, str): self.ans2label = mmengine.load(ans2label) else: self.ans2label = ans2label if self.task == 'vqa' and self.ans2label is not None: self.constraint_trie = Trie(eos=self.tokenizer.eos_token_id) answers = [f' {answer}' for answer in self.ans2label] answer_tokens = self.tokenizer(answers, padding=False) for answer_token in answer_tokens['input_ids']: self.constraint_trie.insert(answer_token) else: self.constraint_trie = None generation_cfg = { 'num_beams': 5, 'max_new_tokens': 20, 'no_repeat_ngram_size': 3, **generation_cfg, } self.model = OFAEncoderDecoder( encoder_cfg=encoder_cfg, decoder_cfg=decoder_cfg, padding_idx=self.tokenizer.pad_token_id, vocab_size=vocab_size, embedding_dim=embedding_dim, generation_cfg=generation_cfg, )
[文档] def forward( self, images: torch.Tensor, data_samples: Optional[list] = None, mode: str = 'predict', **kwargs, ): """The unified entry for a forward process in both training and test. The method accepts the following modes: - "predict": Forward and return a list of data samples contain the predict results. Args: images (torch.Tensor): the preprocessed image tensor of shape ``(N, C, H, W)``. data_samples (List[DataSample], optional): The annotation data of every samples. Defaults to None. mode (str): Return what kind of value. Defaults to 'predict'. """ if mode == 'predict': return self.predict(images, data_samples, **kwargs) else: raise RuntimeError(f'Invalid mode "{mode}".')
def predict( self, images, data_samples=None, post_process=True, **generation_config, ): text_tokens = self.preprocess_text(data_samples, images.size(0), images.device) if 'images_mask' in data_samples[0]: images_mask = torch.tensor([ sample.get('images_mask') for sample in data_samples ]).bool().to(images.device) else: images_mask = None num_beams = generation_config.get( 'num_beams', getattr(self.model.generation_config, 'num_beams')) decoder_prompts = self.get_decoder_prompts(data_samples) constrain_fn = partial( apply_constraint, constraint_trie=self.constraint_trie, decoder_prompts=decoder_prompts, num_beams=num_beams, ) outputs = self.model.generate( input_ids=text_tokens, images=images, images_mask=images_mask, constrain_fn=constrain_fn, **generation_config, ) if decoder_prompts is not None: # Remove the prefix decoder prompt. for prompt_ids, token in zip(decoder_prompts, outputs): token[1:len(prompt_ids) + 1] = self.tokenizer.pad_token_id if post_process: return self.post_process(outputs, data_samples) else: return outputs def get_decoder_prompts(self, data_samples): decoder_prompts = [] if 'decoder_prompt' not in data_samples[0]: return None for sample in data_samples: prompt = ' ' + sample.get('decoder_prompt') prompt_ids = self.tokenizer(prompt, add_special_tokens=False) prompt_ids = prompt_ids['input_ids'] decoder_prompts.append(prompt_ids) return decoder_prompts def preprocess_text(self, data_samples, batch_size, device): if self.task == 'caption': prompt = self.prompt or ' what does the image describe?' prompts = [prompt] * batch_size prompts = self.tokenizer(prompts, return_tensors='pt') return elif self.task == 'vqa': prompts = [] for sample in data_samples: assert 'question' in sample prompt = ' ' + sample.get('question') prompts.append(prompt) prompts = self.tokenizer( prompts, return_tensors='pt', padding=True) return elif self.task == 'refcoco': prompt_template = self.prompt or \ ' which region does the text " {} " describe?' prompts = [] for sample in data_samples: assert 'text' in sample prompt = prompt_template.format(sample.get('text')) prompts.append(prompt) prompts = self.tokenizer( prompts, return_tensors='pt', padding=True) return def post_process(self, outputs, data_samples): out_data_samples = [] if data_samples is None: data_samples = [None] * outputs.size(0) for data_sample, token in zip(data_samples, outputs): if data_sample is None: data_sample = DataSample() if self.task == 'caption': text = self.tokenizer.decode(token, skip_special_tokens=True) text = CleanCaption( lowercase=False, remove_chars=string.punctuation).clean(text) data_sample.pred_caption = text elif self.task == 'vqa': text = self.tokenizer.decode(token, skip_special_tokens=True) data_sample.pred_answer = text.strip() elif self.task == 'refcoco': bbox = token[1:5] - self.tokenizer.bin_offset # During training, the bbox is normalized by 512. It's related # to the `max_image_size` config in the official repo. bbox = bbox / self.tokenizer.num_bins * 512 scale_factor = data_sample.get('scale_factor', (1, 1)) bbox[0::2] /= scale_factor[0] bbox[1::2] /= scale_factor[1] data_sample.pred_bboxes = bbox.unsqueeze(0) if 'gt_bboxes' in data_sample: gt_bboxes = bbox.new_tensor(data_sample.gt_bboxes) gt_bboxes[:, 0::2] /= scale_factor[0] gt_bboxes[:, 1::2] /= scale_factor[1] data_sample.gt_bboxes = gt_bboxes out_data_samples.append(data_sample) return out_data_samples
Read the Docs v: latest
On Read the Docs
Project Home

Free document hosting provided by Read the Docs.