hyavatar/text_encoder/__init__.py

from dataclasses import dataclass
from typing import Optional, Tuple
from copy import deepcopy

import torch
import torch.nn as nn
from transformers import (
    T5EncoderModel, T5Tokenizer,
    CLIPTextModel, CLIPTokenizer, LlavaForConditionalGeneration,
    LlamaModel, LlamaForCausalLM, LlamaTokenizerFast,
    AutoModelForCausalLM, AutoTokenizer,
)
from transformers.utils import ModelOutput

from ..constants import TEXT_ENCODER_PATH, TOKENIZER_PATH
from ..utils.torch_utils import PRECISION_TO_TYPE


def use_default(value, default):
    return value if value is not None else default


def load_text_encoder(text_encoder_type,
                      text_encoder_precision=None,
                      text_encoder_path=None,
                      infer_mode="encoder",
                      logger=None,
                      device=None
                      ):
    if text_encoder_path is None:
        text_encoder_path = TEXT_ENCODER_PATH[text_encoder_type]
    if logger is not None:
        logger.info(f"Loading text encoder model ({text_encoder_type}) from: {text_encoder_path}")

    if text_encoder_type in ["t5", "t5_v11_xxl"]:
        text_encoder = T5EncoderModel.from_pretrained(text_encoder_path)
        text_encoder.final_layer_norm = text_encoder.encoder.final_layer_norm
    elif text_encoder_type == "clipL":
        text_encoder = CLIPTextModel.from_pretrained(text_encoder_path)
        text_encoder.final_layer_norm = text_encoder.text_model.final_layer_norm
    elif text_encoder_type == "llava-llama-3-8b":
        if infer_mode == "encoder":
            text_encoder = LlamaModel.from_pretrained(text_encoder_path, low_cpu_mem_usage=True)
            text_encoder.final_layer_norm = text_encoder.norm
        else:
            text_encoder = LlamaForCausalLM.from_pretrained(text_encoder_path, low_cpu_mem_usage=True)
            text_encoder.final_layer_norm = text_encoder.model.norm
    elif text_encoder_type == "llava-llama-3-8b-vision":
        if infer_mode == "encoder":
            text_encoder = LlavaForConditionalGeneration.from_pretrained(text_encoder_path, low_cpu_mem_usage=True)
            text_encoder.final_layer_norm = text_encoder.language_model.model.norm
        else:
            raise ValueError
    elif text_encoder_type == "glm-4v-9b":
        text_encoder = AutoModelForCausalLM.from_pretrained(
            text_encoder_path, low_cpu_mem_usage=True, trust_remote_code=True)
    else:
        raise ValueError(f"Unsupported text encoder type: {text_encoder_type}")
    # from_pretrained will ensure that the model is in eval mode.

    if text_encoder_precision is not None:
        text_encoder = text_encoder.to(dtype=PRECISION_TO_TYPE[text_encoder_precision])

    text_encoder.requires_grad_(False)

    if logger is not None:
        logger.info(f"Text encoder to dtype: {text_encoder.dtype}")

    if device is not None:
        text_encoder = text_encoder.to(device)

    return text_encoder, text_encoder_path


def load_tokenizer(tokenizer_type,
                   tokenizer_path=None,
                   padding_side="right",
                   logger=None
                   ):
    if tokenizer_path is None:
        tokenizer_path = TOKENIZER_PATH[tokenizer_type]
    if logger is not None:
        logger.info(f"Loading tokenizer ({tokenizer_type}) from: {tokenizer_path}")

    if tokenizer_type in ["t5", "t5_v11_xxl"]:
        tokenizer = T5Tokenizer.from_pretrained(tokenizer_path)
    elif tokenizer_type == "clipL":
        tokenizer = CLIPTokenizer.from_pretrained(tokenizer_path, max_length=77)
    elif tokenizer_type == "llava-llama-3-8b":
        tokenizer = LlamaTokenizerFast.from_pretrained(tokenizer_path, padding_side=padding_side)
    elif tokenizer_type == "glm-4v-9b":
        tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, padding_side=padding_side, trust_remote_code=True)
    else:
        raise ValueError(f"Unsupported tokenizer type: {tokenizer_type}")

    return tokenizer, tokenizer_path


@dataclass
class TextEncoderModelOutput(ModelOutput):
    """
    Base class for model's outputs that also contains a pooling of the last hidden states.

    Args:
        hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
            Sequence of hidden-states at the output of the last layer of the model.
        attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
        hidden_states_list (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed):
            Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
            one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`.
            Hidden-states of the model at the output of each layer plus the optional initial embedding outputs.
        text_outputs (`list`, *optional*, returned when `return_texts=True` is passed):
            List of decoded texts.
    """

    hidden_state: torch.FloatTensor = None
    attention_mask: Optional[torch.LongTensor] = None
    hidden_states_list: Optional[Tuple[torch.FloatTensor, ...]] = None
    text_outputs: Optional[list] = None


class TextEncoder(nn.Module):
    def __init__(self,
                 text_encoder_type: str,
                 max_length: int,
                 text_encoder_precision: Optional[str] = None,
                 text_encoder_path: Optional[str] = None,
                 tokenizer_type: Optional[str] = None,
                 tokenizer_path: Optional[str] = None,
                 output_key: Optional[str] = None,
                 use_attention_mask: bool = True,
                 infer_mode: str = "encoder",
                 input_max_length: Optional[int] = None,
                 prompt_template: Optional[dict] = None,
                 prompt_template_video: Optional[dict] = None,
                 hidden_state_skip_layer: Optional[int] = None,
                 apply_final_norm: bool = False,
                 reproduce: bool = False,
                 logger=None,
                 device=None,
                 ):
        super().__init__()
        self.text_encoder_type = text_encoder_type
        self.max_length = max_length
        self.precision = text_encoder_precision
        self.model_path = text_encoder_path
        self.tokenizer_type = tokenizer_type if tokenizer_type is not None else text_encoder_type
        self.tokenizer_path = tokenizer_path if tokenizer_path is not None else text_encoder_path
        self.use_attention_mask = use_attention_mask
        if prompt_template_video is not None: 
            assert use_attention_mask is True, "Attention mask is True required when training videos."
        self.input_max_length = input_max_length if input_max_length is not None else max_length
        self.prompt_template = prompt_template
        self.prompt_template_video = prompt_template_video
        self.hidden_state_skip_layer = hidden_state_skip_layer
        self.apply_final_norm = apply_final_norm
        self.infer_mode = infer_mode
        self.reproduce = reproduce
        self.logger = logger

        self.use_template = self.prompt_template is not None
        if self.use_template:
            assert isinstance(self.prompt_template, dict) and "template" in self.prompt_template, (
                f"`prompt_template` must be a dictionary with a key 'template', got {self.prompt_template}"
            )
            assert '{}' in str(self.prompt_template["template"]), (
                "`prompt_template['template']` must contain a placeholder `{}` for the input text, "
                f"got {self.prompt_template['template']}"
            )

        self.use_video_template = self.prompt_template_video is not None
        if self.use_video_template:
            if self.prompt_template_video is not None:
                assert isinstance(self.prompt_template_video, dict) and "template" in self.prompt_template_video, (
                    f"`prompt_template_video` must be a dictionary with a key 'template', got {self.prompt_template_video}"
                )
            assert '{}' in str(self.prompt_template_video["template"]), (
                "`prompt_template_video['template']` must contain a placeholder `{}` for the input text, "
                f"got {self.prompt_template_video['template']}"
            )

        if infer_mode == "decoder":
            # TODO(jarvizhang): Move the valid model list to constants.py
            assert text_encoder_type in ["llava-llama-3-8b"], (
                f"Unsupported text encoder type for infer_mode='decoder': {text_encoder_type}"
            )
            assert self.prompt_template is not None and hidden_state_skip_layer is not None, (
                f"`prompt_template` and `hidden_state_skip_layer` must be provided for infer_mode='decoder', "
                f"got prompt_template={self.prompt_template}, hidden_state_skip_layer={self.hidden_state_skip_layer}"
            )

        if "t5" in text_encoder_type:
            self.output_key = output_key or "last_hidden_state"
        elif "clip" in text_encoder_type:
            self.output_key = output_key or "pooler_output"
        elif "llama" in text_encoder_type or "glm" in text_encoder_type:
            self.output_key = output_key or ("last_hidden_state" if infer_mode == "encoder" else None)
        else:
            raise ValueError(f"Unsupported text encoder type: {text_encoder_type}")

        self.model, self.model_path = load_text_encoder(
            text_encoder_type=self.text_encoder_type,
            text_encoder_precision=self.precision,
            text_encoder_path=self.model_path,
            infer_mode=self.infer_mode,
            logger=self.logger,
            device=device
        )
        self.dtype = self.model.dtype
        self.device = self.model.device

        padding_side = "right" if self.infer_mode == "encoder" else "left"
        self.tokenizer, self.tokenizer_path = load_tokenizer(
            tokenizer_type=self.tokenizer_type,
            tokenizer_path=self.tokenizer_path,
            padding_side=padding_side,
            logger=self.logger
        )

    def __repr__(self):
        return f"{self.text_encoder_type} ({self.precision} - {self.model_path})"

    @staticmethod
    def apply_text_to_template(text, template, prevent_empty_text=True):
        """
        Apply text to template.

        Args:
            text (str): Input text.
            template (str or list): Template string or list of chat conversation.
            prevent_empty_text (bool): If Ture, we will prevent the user text from being empty
                by adding a space. Defaults to True.
        """
        if isinstance(template, str):
            # Will send string to tokenizer. Used for llava-llama-3-8b.
            return template.format(text)
        elif isinstance(template, list):
            # Will send chat conversation to tokenizer. Used for glm-4v-9b.
            conversation = deepcopy(template)
            for message_id in range(len(conversation)):
                if '{}' in conversation[message_id]["content"]:
                    filled_text = conversation[message_id]["content"].format(text)
                    if prevent_empty_text and len(filled_text) == 0:
                        filled_text = ' '
                    conversation[message_id]["content"] = filled_text
                    # We assume there is only one placeholder in each conversation.
                    break
            return conversation
        else:
            raise TypeError(f"Unsupported template type: {type(template)}")

    def text2tokens(self, text, data_type='image', name='person'):
        """
        Tokenize the input text.

        Args:
            text (str or list): Input text.
        """
        tokenize_input_type = 'str'
        if self.use_template:
            if data_type == 'image': 
                prompt_template = self.prompt_template["template"]
            elif data_type == 'video': 
                prompt_template = self.prompt_template_video["template"]
            else: 
                raise ValueError(f"Unsupported data type: {data_type}")
            if isinstance(text, (list, tuple)):
                text = [self.apply_text_to_template(one_text, prompt_template) for one_text in text]
                if isinstance(text[0], list):
                    tokenize_input_type = 'list'
            elif isinstance(text, str):
                text = self.apply_text_to_template(text, prompt_template)
                if isinstance(text, list):
                    tokenize_input_type = 'list'
            else:
                raise TypeError(f"Unsupported text type: {type(text)}")

        kwargs = dict(truncation=True, max_length=self.max_length, padding="max_length", return_tensors="pt")
        if self.text_encoder_type == "llava-llama-3-8b-vision":
            if isinstance(text, list):
                for i in range(len(text)):
                    text[i] = text[i] + '\nThe person looks like<image> and the object looks lime<image>'
            elif isinstance(text, str):
                text = text + '\nThe person looks like<image> and the object looks lime<image>'
                # text = text + '\nThe person looks like<image>'
            else:
                raise NotImplementedError
        # print(f'kwargs:{kwargs}')
        # breakpoint()
        if tokenize_input_type == 'str':
            return self.tokenizer(text, return_length=False, return_overflowing_tokens=False, return_attention_mask=True, **kwargs, )
        elif tokenize_input_type == 'list':
            return self.tokenizer.apply_chat_template(text,
                                                      add_generation_prompt=True,
                                                      tokenize=True,
                                                      return_dict=True,
                                                      **kwargs,
                                                      )
        else:
            raise ValueError(f"Unsupported tokenize_input_type: {tokenize_input_type}")

    def encode(self, batch_encoding, use_attention_mask=None, output_hidden_states=False, do_sample=None,
               hidden_state_skip_layer=None, return_texts=False, data_type='image'):
        """
        Args:
            batch_encoding (dict): Batch encoding from tokenizer.
            use_attention_mask (bool): Whether to use attention mask. If None, use self.use_attention_mask.
                Defaults to None.
            output_hidden_states (bool): Whether to output hidden states. If False, return the value of
                self.output_key. If True, return the entire output. If set self.hidden_state_skip_layer,
                output_hidden_states will be set True. Defaults to False.
            do_sample (bool): Whether to sample from the model. Used for Decoder-Only LLMs. Defaults to None.
                When self.produce is False, do_sample is set to True by default.
            hidden_state_skip_layer (int): Number of hidden states to hidden_state_skip_layer. 0 means the last layer.
                If None, self.output_key will be used. Defaults to None.
            return_texts (bool): Whether to return the decoded texts. Defaults to False.
        """
        use_attention_mask = use_default(use_attention_mask, self.use_attention_mask)
        hidden_state_skip_layer = use_default(hidden_state_skip_layer, self.hidden_state_skip_layer)
        do_sample = use_default(do_sample, not self.reproduce)

        if self.infer_mode == "encoder":
            attention_mask = batch_encoding["attention_mask"].to(self.model.device) if use_attention_mask else None
            if 'pixel_value_llava' in batch_encoding:
                outputs = self.model(
                    input_ids=batch_encoding["input_ids"].to(self.model.device),
                    attention_mask=attention_mask,
                    pixel_values=batch_encoding["pixel_value_llava"].to(self.model.device),
                    output_hidden_states=output_hidden_states or hidden_state_skip_layer is not None)
            else:
                outputs = self.model(
                input_ids=batch_encoding["input_ids"].to(self.model.device),
                attention_mask=attention_mask,
                output_hidden_states=output_hidden_states or hidden_state_skip_layer is not None,)
                # breakpoint()
            if hidden_state_skip_layer is not None:
                last_hidden_state = outputs.hidden_states[-(hidden_state_skip_layer + 1)]
                # Real last hidden state already has layer norm applied. So here we only apply it
                # for intermediate layers.
                if hidden_state_skip_layer > 0 and self.apply_final_norm:
                    last_hidden_state = self.model.final_layer_norm(last_hidden_state)
            else:
                last_hidden_state = outputs[self.output_key]

            # Remove hidden states of instruction tokens, only keep prompt tokens.
            if self.use_template:
                if data_type == 'image': 
                    crop_start = self.prompt_template.get("crop_start", -1)
                elif data_type == 'video': 
                    crop_start = self.prompt_template_video.get("crop_start", -1)
                else: 
                    raise ValueError(f"Unsupported data type: {data_type}")
                if crop_start > 0:
                    last_hidden_state = last_hidden_state[:, crop_start:]
                    attention_mask = attention_mask[:, crop_start:] if use_attention_mask else None

            if output_hidden_states:
                return TextEncoderModelOutput(last_hidden_state, attention_mask, outputs.hidden_states)
            return TextEncoderModelOutput(last_hidden_state, attention_mask)

        elif self.infer_mode == "decoder":
            # Remove leading padding tokens
            input_max_valid_tokens = batch_encoding["attention_mask"].sum(dim=1).max().item()
            if input_max_valid_tokens < batch_encoding["attention_mask"].shape[1]:
                batch_encoding = {
                    "input_ids": batch_encoding["input_ids"][:, -input_max_valid_tokens:],
                    "attention_mask": batch_encoding["attention_mask"][:, -input_max_valid_tokens:],
                }

            # Generate text from the model.
            outputs = self.model.generate(
                input_ids=batch_encoding["input_ids"].to(self.model.device),
                attention_mask=batch_encoding["attention_mask"].to(self.model.device) if use_attention_mask else None,
                max_new_tokens=self.max_length,
                do_sample=do_sample,
                return_dict_in_generate=True,
                output_hidden_states=True,
                stop_strings='<|eot_id|>', tokenizer=self.tokenizer,
                pad_token_id=self.tokenizer.eos_token_id,
            )

            # Concatenate hidden states from all generated tokens.
            hidden_states = torch.cat([
                per_token_hidden_states[-(hidden_state_skip_layer + 1)]
                for per_token_hidden_states in outputs.hidden_states[1:]
            ], dim=1)
            if self.apply_final_norm:
                hidden_states = self.model.final_layer_norm(hidden_states)

            # Make sequence mask from output sequences
            output_max_valid_tokens = hidden_states.shape[1]
            attention_mask = (outputs.sequences[:, -output_max_valid_tokens - 1:-1] != self.tokenizer.eos_token_id).long()

            if return_texts:
                text_outputs = self.tokenizer.batch_decode(outputs.sequences, skip_special_tokens=False)
                return TextEncoderModelOutput(hidden_states, attention_mask, None, text_outputs)
            else:
                return TextEncoderModelOutput(hidden_states, attention_mask)
        else:
            raise ValueError(f"Unsupported text encoder infer mode: {self.infer_mode}")

    def forward(self, text, use_attention_mask=None, output_hidden_states=False, do_sample=False,
                hidden_state_skip_layer=None, return_texts=False):
        batch_encoding = self.text2tokens(text)
        return self.encode(batch_encoding, use_attention_mask=use_attention_mask,
                           output_hidden_states=output_hidden_states, do_sample=do_sample,
                           hidden_state_skip_layer=hidden_state_skip_layer, return_texts=return_texts)