conditioning/t5embedder.py

import logging
from transformers import T5Tokenizer, T5EncoderModel
from typing import List, Optional, Dict, Sequence, Tuple
import random
import torch
from torch import Tensor

from conditioning.embedder import Embedder, LinearProjectionEmbedder
from conditioning.embedded_condition import EmbeddedCondition


class T5Embedder(LinearProjectionEmbedder):
    """T5-based TextConditioner.

    Args:
        name (str): Name of the T5 model.
        output_dim (int): Output dim of the conditioner.
        finetune (bool): Whether to fine-tune T5 at train time.
        word_dropout (float, optional): Word dropout probability.
        normalize_text (bool, optional): Whether to apply text normalization.
    """
    MODELS = [
        "t5-small", "t5-base", "t5-large", "t5-3b", "t5-11b",
        "google/flan-t5-small", "google/flan-t5-base", "google/flan-t5-large",
        "google/flan-t5-xl", "google/flan-t5-xxl"
    ]
    MODELS_DIMS = {
        "t5-small": 512,
        "t5-base": 768,
        "t5-large": 1024,
        "t5-3b": 1024,
        "t5-11b": 1024,
        "google/flan-t5-small": 512,
        "google/flan-t5-base": 768,
        "google/flan-t5-large": 1024,
        "google/flan-t5-3b": 1024,
        "google/flan-t5-11b": 1024,
    }

    def __init__(self,
                 embedding_dim: int,
                 t5_on_cpu: bool,
                 name: str = "t5-base",
                 finetune: bool = False,
                 word_dropout: float = 0.3):
        assert name in self.MODELS, f"Unrecognized t5 model name (should in {self.MODELS})"
        super().__init__(self.MODELS_DIMS[name], embedding_dim)
        self.name = name
        self.finetune = finetune
        self.word_dropout = word_dropout
        self.t5_on_cpu: bool = t5_on_cpu
        if self.t5_on_cpu and finetune:
            raise ValueError("Can't finetune t5 if it's locked on cpu")

        # Let's disable logging temporarily because T5 will vomit some errors otherwise.
        # thanks https://gist.github.com/simon-weber/7853144
        previous_level = logging.root.manager.disable
        logging.disable(logging.ERROR)
        self.t5_tokenizer = T5Tokenizer.from_pretrained(
            name,
            clean_up_tokenization_spaces=False,
        )

        t5 = T5EncoderModel.from_pretrained(name).train(mode=finetune)

        if self.t5_on_cpu:
            t5 = t5.cpu()
            self.__dict__["t5"] = t5
        else:
            self.t5 = t5

        if not self.finetune:
            for p in self.t5.parameters():
                p.requires_grad = False

        # if not self.finetune:
        #     if self.t5_on_cpu:
        #         t5 = t5.cpu()
        #     self.__dict__["t5"] = t5.eval()
        #     for p in self.t5.parameters():
        #         p.requires_grad = False
        # else:
        #     self.t5 = t5

        # with warnings.catch_warnings():

        #     warnings.simplefilter("ignore")
        #     try:
        #         self.t5_tokenizer = T5Tokenizer.from_pretrained(name)
        #         t5 = T5EncoderModel.from_pretrained(name).train(mode=finetune)
        #     finally:
        #         logging.disable(previous_level)
        # if finetune:
        #     self.t5 = t5
        # else:
        #     # this makes sure that the t5 models is not part
        #     # of the saved checkpoint
        #     self.__dict__['t5'] = t5

        # self.normalize_text = normalize_text
        # if normalize_text:
        #     self.text_normalizer = WhiteSpaceTokenizer(1,
        #                                                lemma=True,
        #                                                stopwords=True)

    # def to(self, *args, **kwargs):
    #     return super().to("cpu")

    # def cuda(self, *args, **kwargs):
    #     return self.to("cpu")

    def tokenize(self, x: Sequence[Optional[str]]) -> Dict[str, torch.Tensor]:
        # if current sample doesn't have a certain attribute, replace with empty string
        entries: List[str] = [xi if xi is not None else "" for xi in x]
        # if self.normalize_text:
        #     _, _, entries = self.text_normalizer(  # type: ignore
        #         entries, return_text=True)
        if self.word_dropout > 0. and self.training:
            new_entries = []
            for entry in entries:
                words = [
                    word for word in entry.split(" ")
                    if random.random() >= self.word_dropout
                ]
                new_entries.append(" ".join(words))
            entries = new_entries

        empty_idx = torch.LongTensor(
            [i for i, xi in enumerate(entries) if xi == ""])

        inputs = self.t5_tokenizer(entries, return_tensors='pt',
                                   padding=True).to(
                                       next(iter(self.parameters())).device)
        mask = inputs['attention_mask']
        mask[empty_idx, :] = 0  # zero-out index where the input is non-existant
        return inputs

    def forward(self,
                descriptions: List[str],
                duplicate_for_cfg: bool = False) -> EmbeddedCondition:

        if duplicate_for_cfg:
            descriptions = descriptions + ([""] * len(descriptions))

        tokenized = self.tokenize(descriptions)
        embedding, mask = self.embed(tokenized)
        mask = mask.bool()
        return EmbeddedCondition(embedding, mask)

    def embed(self, inputs: Dict[str, torch.Tensor]) -> Tuple[Tensor, Tensor]:
        if not self.finetune:
            self.t5.eval()
        mask = inputs['attention_mask']
        with torch.set_grad_enabled(self.finetune):
            if self.t5_on_cpu:
                inputs = {k: v.to("cpu") for k, v in inputs.items()}
            with torch.autocast(device_type="cuda", enabled=False):
                embeds = self.t5(**inputs).last_hidden_state
        embeds = self.output_proj(embeds.to(self.output_proj.weight))
        embeds = (embeds * mask.unsqueeze(-1).to(embeds))
        return embeds, mask

    def null_condition(self, batch_size: int):
        return EmbeddedCondition(
            torch.zeros(batch_size,
                        1,
                        self.embedding_dim,
                        dtype=torch.float32,
                        device=self.output_proj.weight.device),
            torch.ones(batch_size,
                       1,
                       dtype=torch.bool,
                       device=self.output_proj.weight.device))


class T5EmbedderCPU(T5Embedder):

    def __init__(self, embedding_dim: int):
        super().__init__(embedding_dim,
                         t5_on_cpu=True,
                         name="t5-base",
                         finetune=False,
                         word_dropout=0.3)


class T5EmbedderGPU(T5Embedder):

    def __init__(self, embedding_dim: int):
        super().__init__(embedding_dim,
                         t5_on_cpu=False,
                         name="t5-base",
                         finetune=False,
                         word_dropout=0.3)


if __name__ == "__main__":
    from utils.inspection import print_params
    from time import time

    t5 = T5EmbedderCPU(1024).eval().cpu()
    # print_params(t5, 1, False)

    descriptions = ["the quick"]

    embedded_desc = t5(descriptions)
    # print(embedded_desc)
    # print(embedded_desc.data)
    # print(embedded_desc.mask)

    # print(embedded_desc.data[..., -1, ...])
    # print(embedded_desc.data[0, -1, ...])