MultimodalOCR/models/mPLUG_Owl/mplug_owl/processing_mplug_owl.py

import re

import torch
import torch.utils.checkpoint

from transformers.processing_utils import ProcessorMixin
from transformers.tokenization_utils_base import BatchEncoding
from transformers.models.clip.image_processing_clip import CLIPImageProcessor


media_token = ({"image": ("<image>", 65)},)


class MplugOwlProcessor(ProcessorMixin):
    attributes = []
    tokenizer_class = ("MplugOwlTokenizer")

    def __init__(self, image_processor=None, tokenizer=None, **kwargs):
        super().__init__(**kwargs)
        self.tokens_to_generate = 0
        self.image_processor = image_processor
        self.tokenizer = tokenizer
        self.add_BOS = True

    def __call__(self, text=None, images=None, return_tensors=None, **kwargs):
        if text is None and images is None:
            raise ValueError("You have to specify either text or images. Both cannot be none.")

        if text is not None:
            encoding = tokenize_prompts(
                prompts=text,
                tokens_to_generate=self.tokens_to_generate,
                add_BOS=self.add_BOS,
                tokenizer=self.tokenizer,
                ignore_dist=True,
                **kwargs,
            )
            # encoding = self.tokenizer(text, return_tensors=return_tensors, **kwargs)

        if images is not None:
            image_features = self.image_processor(images, return_tensors=return_tensors, **kwargs)

        if text is not None and images is not None:
            encoding["pixel_values"] = image_features.pixel_values
            return BatchEncoding(data=encoding)
        elif text is not None:
            return BatchEncoding(data=encoding)
        else:
            return BatchEncoding(data=dict(**image_features), tensor_type=return_tensors)

    def batch_decode(self, skip_special_tokens=True, *args, **kwargs):
        """
        This method forwards all its arguments to CLIPTokenizerFast's [`~PreTrainedTokenizer.batch_decode`]. Please
        refer to the docstring of this method for more information.
        """
        return self.tokenizer.batch_decode(*args, skip_special_tokens=skip_special_tokens, **kwargs)

    def decode(self, skip_special_tokens=True, *args, **kwargs):
        """
        This method forwards all its arguments to CLIPTokenizerFast's [`~PreTrainedTokenizer.decode`]. Please refer to
        the docstring of this method for more information.
        """
        return self.tokenizer.decode(*args, skip_special_tokens=skip_special_tokens, **kwargs)


class MplugOwlImageProcessor(CLIPImageProcessor):
    pass


def detokenize_generations(tokens_gpu_tensor, lengths_gpu_tensor, return_segments, tokenizer):
    """Detokenize the generated tokens."""

    prompts_plus_generations = []
    if return_segments:
        prompts_plus_generations_segments = []

    tokens = tokens_gpu_tensor.cpu().numpy().tolist()
    lengths = lengths_gpu_tensor.cpu().numpy().tolist()
    for sequence_tokens, length in zip(tokens, lengths):
        sequence_tokens = sequence_tokens[:length]
        prompts_plus_generations.append(tokenizer.detokenize(sequence_tokens))
        if return_segments:
            from tokenizers.decoders import Metaspace

            if hasattr(tokenizer, "tokenizer"):
                if isinstance(tokenizer.tokenizer.decoder, Metaspace):
                    words = tokenizer.tokenizer.decode(sequence_tokens)
                else:
                    words = []
                    for token in sequence_tokens:
                        word = tokenizer.tokenizer.decoder[token]
                        word = bytearray([tokenizer.tokenizer.byte_decoder[c] for c in word]).decode(
                            "utf-8", errors="replace"
                        )
                        words.append(word)
                prompts_plus_generations_segments.append(words)
            else:
                words = tokenizer.detokenize(sequence_tokens)
                # else:
                #     words = []
                #     for token in sequence_tokens:
                #         word = tokenizer.tokenizer.decoder[token]
                #         word = bytearray(
                #             [tokenizer.tokenizer.byte_decoder[c] for c in word]).decode(
                #                 'utf-8', errors='replace')
                #         words.append(word)
                prompts_plus_generations_segments.append(words)

    if return_segments:
        return tokens, prompts_plus_generations, prompts_plus_generations_segments

    return tokens, prompts_plus_generations


def tokenize_prompts(
    prompts=None, tokens_to_generate=None, add_BOS=None, rank=0, tokenizer=None, ignore_dist=False, **kwargs
):
    """Tokenize prompts and make them avaiable on all ranks."""

    # On all ranks set to None so we can pass them to functions
    prompts_tokens_cuda_long_tensor = None
    prompts_length_cuda_long_tensor = None

    # On the specified rank, build the above.
    attention_mask = None
    if ignore_dist or torch.distributed.get_rank() == rank:
        assert prompts is not None
        assert tokens_to_generate is not None
        # Tensor of tokens padded and their unpadded length.
        prompts_tokens_cuda_long_tensor, prompts_length_cuda_long_tensor, attention_mask = _tokenize_prompts_and_batch(
            prompts, tokens_to_generate, add_BOS, tokenizer, **kwargs
        )

    return {
        "input_ids": prompts_tokens_cuda_long_tensor,
        "attention_mask": attention_mask,
        # "prompt_length": prompts_length_cuda_long_tensor,
    }


def _tokenize_prompts_and_batch(prompts, tokens_to_generate, add_BOS, tokenizer, **kwargs):
    """Given a set of prompts and number of tokens to generate:
    - tokenize prompts
    - set the sequence length to be the max of length of prompts plus the number of tokens we would like to generate
    - pad all the sequences to this length so we can convert them into a 2D tensor.
    """

    # Tokenize all the prompts.
    # if add_BOS:
    #     prompts_tokens = [[tokenizer.bos] + tokenizer.tokenize(prompt)
    #                       for prompt in prompts]
    # else:
    #     prompts_tokens = [tokenizer.tokenize(prompt) for prompt in prompts]

    prompts_tokens = [_tokenize_prompt(prompt, tokenizer, add_BOS, **kwargs) for prompt in prompts]

    # Now we have a list of list of tokens which each list has a different
    # size. We want to extend this list to:
    #   - incorporate the tokens that need to be generated
    #   - make all the sequences equal length.
    # Get the prompts length.
    prompts_length = [len(prompt_tokens) for prompt_tokens in prompts_tokens]
    # Get the max prompts length.
    max_prompt_len = max(prompts_length)
    # Number of tokens in the each sample of the batch.
    samples_length = max_prompt_len + tokens_to_generate
    # Now update the list of list to be of the same size: samples_length.
    for prompt_tokens, prompt_length in zip(prompts_tokens, prompts_length):
        padding_size = samples_length - prompt_length
        prompt_tokens.extend([tokenizer.eos_token_id] * padding_size)

    # Now we are in a structured format, we can convert to tensors.
    prompts_tokens_tensor = torch.LongTensor(prompts_tokens)
    prompts_length_tensor = torch.LongTensor(prompts_length)
    attention_mask = torch.zeros(prompts_tokens_tensor.shape[:2])
    for i, l in enumerate(prompts_length_tensor):
        attention_mask[i, :l] = 1
    return prompts_tokens_tensor, prompts_length_tensor, attention_mask


def _tokenize_prompt(prompt, tokenizer, add_BOS=False, media_info={"<image>": 65}, **kwargs):
    media_tokens = {k: -int(i + 1) for i, k in enumerate(media_info.keys())}
    media_lengths = media_info.copy()

    if add_BOS:
        prompt_chunk = [tokenizer.bos_token_id]
    else:
        prompt_chunk = []

    # Pure Text
    if all([media_token not in prompt for media_token in media_tokens.keys()]):
        enc_chunk = prompt_chunk + tokenizer(prompt, add_special_tokens=False, **kwargs)["input_ids"]

    # Multi-Modal Text
    else:
        enc_chunk = prompt_chunk
        pattern = "|".join(map(re.escape, list(media_tokens.keys())))
        chunk_strs = re.split(f"({pattern})", prompt)
        chunk_strs = [x for x in chunk_strs if len(x) > 0]
        for idx, chunk_str in enumerate(chunk_strs):
            if chunk_str in media_tokens:
                enc_chunk += [media_tokens[chunk_str]] * media_lengths[chunk_str]
            else:
                tmp_chunk = tokenizer(chunk_str, add_special_tokens=False)["input_ids"]
                # if idx < len(chunk_strs) - 1: # Last chunk should not have eos
                #     tmp_chunk += [tokenizer.eod_id]
                enc_chunk += tmp_chunk
    return enc_chunk