yeelou/design2code-hf-bit8

1"""largely copy from llama and adapt for CogAgent"""2 3import math4import warnings5from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union6 7import torch8from einops import rearrange9from torch import nn10 11# from .util import FastRotaryEmbedding12from torch.nn import CrossEntropyLoss13from torch.nn import functional as F14from torchvision import transforms15from transformers import PreTrainedModel, PreTrainedTokenizer16from transformers.activations import ACT2FN17from transformers.modeling_outputs import (18    BaseModelOutputWithPast,19    CausalLMOutputWithPast,20)21from transformers.utils.logging import get_logger22 23from .configuration_cogagent import CogAgentConfig24from .cross_visual import CrossVisionModel25from .visual import EVA2CLIPModel26 27if TYPE_CHECKING:28    from transformers.utils import ModelOutput29 30logger = get_logger(__name__)31 32LANGUAGE_TOKEN_TYPE = 033VISION_TOKEN_TYPE = 134 35 36# Copied from transformers.models.bart.modeling_bart._make_causal_mask37def _make_causal_mask(38    input_ids_shape: torch.Size,39    dtype: torch.dtype,40    device: torch.device,41    past_key_values_length: int = 0,42):43    """44    Make causal mask used for bi-directional self-attention.45    """46    bsz, tgt_len = input_ids_shape47    mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)48    mask_cond = torch.arange(mask.size(-1), device=device)49    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)50    mask = mask.to(dtype)51 52    if past_key_values_length > 0:53        mask = torch.cat(54            [55                torch.zeros(56                    tgt_len, past_key_values_length, dtype=dtype, device=device57                ),58                mask,59            ],60            dim=-1,61        )62    return mask[None, None, :, :].expand(63        bsz, 1, tgt_len, tgt_len + past_key_values_length64    )65 66 67# Copied from transformers.models.bart.modeling_bart._expand_mask68def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):69    """70    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.71    """72    bsz, src_len = mask.size()73    tgt_len = tgt_len if tgt_len is not None else src_len74 75    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)76 77    inverted_mask = 1.0 - expanded_mask78 79    return inverted_mask.masked_fill(80        inverted_mask.to(torch.bool), torch.finfo(dtype).min81    )82 83 84class RMSNorm(nn.Module):85    def __init__(self, hidden_size, eps=1e-6):86        super().__init__()87        self.weight = nn.Parameter(torch.ones(hidden_size))88        self.variance_epsilon = eps89 90    def forward(self, hidden_states):91        input_dtype = hidden_states.dtype92        hidden_states = hidden_states.to(torch.float32)93        variance = hidden_states.pow(2).mean(-1, keepdim=True)94        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)95        return (self.weight * hidden_states).to(input_dtype)96 97 98class MLP(nn.Module):99    def __init__(self, config):100        super().__init__()101        self.hidden_size = config.hidden_size102        self.intermediate_size = config.intermediate_size103        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)104        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)105        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)106        self.act_fn = ACT2FN[config.hidden_act]107 108    def forward(self, x):109        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))110        return down_proj111 112 113def get_expert_mask(114    token_type_ids: "torch.LongTensor(B, L)",115) -> "[torch.BoolTensor(B, L), torch.BoolTensor(B, L)]":116    vision_token_mask = torch.zeros_like(token_type_ids, dtype=torch.bool)117    vision_token_mask[:, :-1] = (token_type_ids[:, :-1] == VISION_TOKEN_TYPE) & (118        token_type_ids[:, 1:] == VISION_TOKEN_TYPE119    )120    language_token_mask = ~vision_token_mask121    return vision_token_mask, language_token_mask122 123 124class VisionExpertMLP(nn.Module):125    def __init__(self, config):126        super().__init__()127        self.language_mlp = MLP(config)128        self.vision_mlp = MLP(config)129 130    def forward(131        self,132        hidden_states: "torch.Tensor(B, L, D)",133        token_type_ids: "torch.LongTensor(B, L)",134    ):135        output = torch.empty(136            hidden_states.shape, dtype=hidden_states.dtype, device=hidden_states.device137        )138        vision_token_mask, language_token_mask = get_expert_mask(token_type_ids)139        output[vision_token_mask] = self.vision_mlp(hidden_states[vision_token_mask])140        output[language_token_mask] = self.language_mlp(141            hidden_states[language_token_mask]142        )143        return output144 145 146def attention_fn(147    query_layer: "torch.tensor(B, H, L, HD)",148    key_layer: "torch.tensor(B, H, L, HD)",149    value_layer: "torch.tensor(B, H, L, HD)",150    attention_mask: "torch.tensor(B, H, L, HD)",151    *,152    scaling_attention_score: bool = True,153    attention_dropout: nn.Module = None,154):155    attention_mask_bool = attention_mask == 0156    is_low_triangle = (157        attention_mask_bool158        == torch.ones_like(attention_mask_bool, dtype=torch.float).tril()159    ).all()160    is_full = (attention_mask_bool > 0).all()161    if not (int(torch.__version__.split(".")[0]) >= 2):162        warnings.warn("It's recommended to use torch2.0 or higher.")163    if (164        int(torch.__version__.split(".")[0]) >= 2165        and scaling_attention_score166        and (is_full or is_low_triangle)167    ):168        dropout_p = (169            0.0170            if attention_dropout is None or not attention_dropout.training171            else attention_dropout.p172        )173        return torch.nn.functional.scaled_dot_product_attention(174            query_layer,175            key_layer,176            value_layer,177            attn_mask=None,178            dropout_p=dropout_p,179            is_causal=not is_full,180        )181    else:182        if scaling_attention_score:183            query_layer = query_layer / math.sqrt(query_layer.shape[-1])184        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))185        attention_scores = attention_scores + attention_mask186        attention_scores = nn.functional.softmax(187            attention_scores, dim=-1, dtype=torch.float32188        ).to(query_layer.dtype)189        if attention_dropout is not None:190            attention_scores = attention_dropout(attention_scores)191        context_layer = torch.matmul(attention_scores, value_layer)192        return context_layer193 194 195class RotaryEmbedding(torch.nn.Module):196    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):197        super().__init__()198 199        self.dim = dim200        self.max_position_embeddings = max_position_embeddings201        self.base = base202        inv_freq = self._compute_inv_freq(device)203        self.register_buffer("inv_freq", inv_freq)204        self.max_seq_len_cached = 0205 206    def _compute_inv_freq(self, device=None):207        return 1.0 / (208            self.base ** (torch.arange(0, self.dim, 2, device=device) / self.dim)209        )210 211    def _set_cos_sin_cache(self, seq_len, device, dtype):212        self.max_seq_len_cached = seq_len213        t = torch.arange(214            self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype215        )216 217        freqs = torch.einsum("i,j->ij", t, self.inv_freq)218        # Different from paper, but it uses a different permutation in order to obtain the same calculation219        emb = torch.cat((freqs, freqs), dim=-1)220        self.register_buffer(221            "cos_cached", emb.cos()[:, None, :].to(dtype), persistent=False222        )223        self.register_buffer(224            "sin_cached", emb.sin()[:, None, :].to(dtype), persistent=False225        )226 227    def forward(self, x, seq_len):228        # x: [bs, num_attention_heads, seq_len, head_size]229        if seq_len > self.max_seq_len_cached:230            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)231 232        return (233            self.cos_cached[:seq_len, ...].to(dtype=x.dtype),234            self.sin_cached[:seq_len, ...].to(dtype=x.dtype),235        )236 237 238def rotate_half(x):239    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]240    return torch.cat((-x2, x1), dim=x1.ndim - 1)241 242 243def apply_rotary_pos_emb_index_bhs(q, k, cos, sin, position_id):244    # batch_size, num_head, seq_len, hidden_size245    cos, sin = F.embedding(position_id, cos.squeeze(1)).unsqueeze(1), F.embedding(246        position_id, sin.squeeze(1)247    ).unsqueeze(1)248    q, k = (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)249    return q, k250 251 252class VisionExpertAttention(nn.Module):253    def __init__(self, config):254        super().__init__()255        self.config = config256        self.hidden_size = config.hidden_size257        self.num_heads = config.num_attention_heads258        self.head_dim = self.hidden_size // self.num_heads259        self.max_position_embeddings = config.max_position_embeddings260 261        self.rotary_emb = RotaryEmbedding(self.head_dim)262        self.vision_expert_query_key_value = nn.Linear(263            self.hidden_size, self.hidden_size * 3, bias=False264        )265        self.vision_expert_dense = nn.Linear(266            self.hidden_size, self.hidden_size, bias=False267        )268        self.language_expert_query_key_value = nn.Linear(269            self.hidden_size, self.hidden_size * 3, bias=False270        )271        self.language_expert_dense = nn.Linear(272            self.hidden_size, self.hidden_size, bias=False273        )274 275    def _transpose_for_scores(self, tensor):276        """Transpose a 3D tensor [B, L, H*HD] into a 4D tensor with size [B H L HD]."""277        new_tensor_shape = tensor.size()[:-1] + (self.num_heads, self.head_dim)278        tensor = tensor.view(*new_tensor_shape)279        return tensor.permute(0, 2, 1, 3)280 281    def forward(282        self,283        hidden_states: torch.Tensor,284        token_type_ids: torch.LongTensor,285        position_ids: torch.LongTensor,286        attention_mask: Optional[torch.Tensor] = None,287        past_key_value: Optional[Tuple[torch.Tensor]] = None,288        output_attentions: bool = False,289        use_cache: bool = False,290    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:291        bsz, q_len, _ = hidden_states.size()292        vision_token_mask, language_token_mask = get_expert_mask(token_type_ids)293 294        shape = list(hidden_states.shape)295        shape[-1] = shape[-1] * 3296        mixed_raw_layer = torch.empty(297            shape, dtype=hidden_states.dtype, device=hidden_states.device298        )299        mixed_raw_layer[vision_token_mask] = self.vision_expert_query_key_value(300            hidden_states[vision_token_mask]301        )302        mixed_raw_layer[language_token_mask] = self.language_expert_query_key_value(303            hidden_states[language_token_mask]304        )305 306        query_states, key_states, value_states = torch.split(307            mixed_raw_layer, self.hidden_size, dim=-1308        )309        query_states = self._transpose_for_scores(query_states)  # B, H, L, HD310        key_states = self._transpose_for_scores(key_states)  # B, H, L, HD311        value_states = self._transpose_for_scores(value_states)  # B, H, L, HD312 313        kv_seq_len = key_states.shape[-2]314        if past_key_value is not None:315            kv_seq_len += past_key_value[0].shape[-2]316 317        cos, sin = self.rotary_emb(value_states, seq_len=position_ids.max() + 1)318        query_states, key_states = apply_rotary_pos_emb_index_bhs(319            query_states, key_states, cos, sin, position_ids320        )321 322        if past_key_value is not None:323            key_states = torch.cat([past_key_value[0], key_states], dim=2)324            value_states = torch.cat([past_key_value[1], value_states], dim=2)325 326        past_key_value = (key_states, value_states) if use_cache else None327 328        context_layer = attention_fn(329            query_layer=query_states,330            key_layer=key_states,331            value_layer=value_states,332            attention_mask=attention_mask,333            scaling_attention_score=True,334            attention_dropout=None,335        )336        if context_layer.size() != (bsz, self.num_heads, q_len, self.head_dim):337            raise ValueError(338                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"339                f" {context_layer.size()}"340            )341        context_layer = (342            context_layer.transpose(1, 2)343            .contiguous()344            .reshape(bsz, q_len, self.hidden_size)345        )346 347        attn_output = torch.empty(348            context_layer.shape, dtype=hidden_states.dtype, device=hidden_states.device349        )350        attn_output[vision_token_mask] = self.vision_expert_dense(351            context_layer[vision_token_mask]352        )353        attn_output[language_token_mask] = self.language_expert_dense(354            context_layer[language_token_mask]355        )356 357        if output_attentions:358            warnings.warn("output_attentions is not implemented.")359 360        return attn_output, None, past_key_value361 362 363class CrossAttention(nn.Module):364    def __init__(self, config):365        super().__init__()366        self.config = config367        self.hidden_size = config.hidden_size368        self.cross_hidden_size = config.cross_hidden_size369        self.cross_compute_hidden_size = config.cross_compute_hidden_size370        self.num_heads = config.num_attention_heads371        self.head_dim = self.hidden_size // self.num_heads372        self.cross_head_dim = self.cross_compute_hidden_size // self.num_heads373        self.max_position_embeddings = config.max_position_embeddings374 375        self.query = nn.Linear(376            self.hidden_size, self.cross_compute_hidden_size, bias=False377        )378        self.key_value = nn.Linear(379            self.cross_hidden_size, self.cross_compute_hidden_size * 2, bias=False380        )381        self.dense = nn.Linear(382            self.cross_compute_hidden_size, self.hidden_size, bias=False383        )384 385    def _transpose_for_scores(self, tensor):386        """Transpose a 3D tensor [B, L, H*HD] into a 4D tensor with size [B H L HD]."""387        new_tensor_shape = tensor.size()[:-1] + (self.num_heads, self.cross_head_dim)388        tensor = tensor.view(*new_tensor_shape)389        return tensor.permute(0, 2, 1, 3)390 391    def forward(392        self,393        hidden_states: torch.Tensor,394        encoder_outputs: torch.LongTensor,395        attention_mask: Optional[torch.Tensor] = None,396        past_key_value: Optional[Tuple[torch.Tensor]] = None,397        output_attentions: bool = False,398        use_cache: bool = False,399    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:400        bsz, q_len, _ = hidden_states.size()401 402        shape = list(hidden_states.shape)403        shape[-1] = shape[-1] * 3404 405        mixed_query_layer = self.query(hidden_states)406        if past_key_value is None:407            mixed_x_layer = self.key_value(encoder_outputs)408            mixed_key_layer, mixed_value_layer = torch.split(409                mixed_x_layer, self.cross_compute_hidden_size, dim=-1410            )411            key_states = self._transpose_for_scores(mixed_key_layer)  # B, H, L, HD412            value_states = self._transpose_for_scores(mixed_value_layer)  # B, H, L, HD413        else:414            key_states, value_states = past_key_value415 416        query_states = self._transpose_for_scores(mixed_query_layer)  # B, H, L, HD417 418        past_key_value = (key_states, value_states) if use_cache else None419 420        context_layer = attention_fn(421            query_layer=query_states,422            key_layer=key_states,423            value_layer=value_states,424            attention_mask=attention_mask,425            scaling_attention_score=True,426            attention_dropout=None,427        )428        if context_layer.size() != (bsz, self.num_heads, q_len, self.cross_head_dim):429            raise ValueError(430                f"`cross_attn_output` should be of size {(bsz, self.num_heads, q_len, self.cross_head_dim)}, but is"431                f" {context_layer.size()}"432            )433        context_layer = (434            context_layer.transpose(1, 2)435            .contiguous()436            .reshape(bsz, q_len, self.cross_hidden_size)437        )438 439        attn_output = self.dense(context_layer)440 441        if output_attentions:442            warnings.warn("output_attentions is not implemented.")443 444        return attn_output, None, past_key_value445 446 447class CogAgentDecoderLayer(nn.Module):448    def __init__(self, config):449        super().__init__()450        self.hidden_size = config.hidden_size451        self.self_attn = VisionExpertAttention(config=config)452        self.cross_attn = CrossAttention(config=config)453        self.mlp = VisionExpertMLP(config)454        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)455        self.post_attention_layernorm = RMSNorm(456            config.hidden_size, eps=config.rms_norm_eps457        )458        self.post_cross_attention_layernorm = RMSNorm(459            config.hidden_size, eps=config.rms_norm_eps460        )461 462    def forward(463        self,464        hidden_states: torch.Tensor,465        encoder_outputs: torch.Tensor,466        token_type_ids: torch.LongTensor,467        position_ids: torch.LongTensor,468        attention_mask: Optional[torch.Tensor] = None,469        cross_attention_mask: Optional[torch.Tensor] = None,470        past_key_value: Optional[Tuple[torch.Tensor]] = None,471        output_attentions: Optional[bool] = False,472        use_cache: Optional[bool] = False,473    ) -> Tuple[474        torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]475    ]:476        residual = hidden_states477 478        hidden_states = self.input_layernorm(hidden_states)479 480        # Self Attention481        hidden_states, self_attn_weights, present_key_value = self.self_attn(482            hidden_states=hidden_states,483            token_type_ids=token_type_ids,484            position_ids=position_ids,485            attention_mask=attention_mask,486            past_key_value=past_key_value[:2] if past_key_value is not None else None,487            output_attentions=output_attentions,488            use_cache=use_cache,489        )490        hidden_states = residual + hidden_states491 492        cross_input = self.post_cross_attention_layernorm(hidden_states)493        # Fully Connected494        attention_output, self_cross_attn_weights, present_cross_key_value = (495            self.cross_attn(496                hidden_states=cross_input,497                encoder_outputs=encoder_outputs,498                attention_mask=cross_attention_mask,499                past_key_value=(500                    past_key_value[-2:] if past_key_value is not None else None501                ),502                output_attentions=output_attentions,503                use_cache=use_cache,504            )505        )506        hidden_states = hidden_states + attention_output507        mlp_input = self.post_attention_layernorm(hidden_states)508        mlp_output = self.mlp(mlp_input, token_type_ids=token_type_ids)509        hidden_states = mlp_output + hidden_states510 511        outputs = (hidden_states,)512 513        if output_attentions:514            outputs += (self_attn_weights,)515 516        if use_cache:517            outputs += (present_key_value + present_cross_key_value,)518 519        return outputs  # type: ignore520 521 522class CogAgentPreTrainedModel(PreTrainedModel):523    config_class = CogAgentConfig524    base_model_prefix = "model"525    supports_gradient_checkpointing = False526    _no_split_modules = ["CogAgentDecoderLayer", "TransformerLayer", "Block"]527    _skip_keys_device_placement = "past_key_values"528 529    def _init_weights(self, module):530        std = self.config.initializer_range531        if isinstance(module, nn.Linear):532            module.weight.data.normal_(mean=0.0, std=std)533            if module.bias is not None:534                module.bias.data.zero_()535        elif isinstance(module, nn.Embedding):536            module.weight.data.normal_(mean=0.0, std=std)537            if module.padding_idx is not None:538                module.weight.data[module.padding_idx].zero_()539 540 541def is_empty(images_list: Optional[List[List[torch.Tensor]]]):542    if images_list is None or len(images_list) == 0:543        return True544    for image_list in images_list:545        if len(image_list):546            return False547    return True548 549 550def build_position_ids(551    x: "torch.BoolTensor(B, L)",552    attention_mask: Optional["torch.BoolTensor(B, L)"] = None,553) -> "torch.LongTensor(B, L)":554    if attention_mask is not None:555        tmp = x.clone()556        tmp[~(attention_mask.bool())] = -1557    else:558        tmp = x.clone()559    # image boi eoi token as LANGUAGE_TOKEN_TYPE560    is_boi_eoi = torch.zeros_like(x, dtype=torch.bool)561    is_boi_eoi[:, 1:] |= (tmp[:, 1:] == VISION_TOKEN_TYPE) & (562        tmp[:, :-1] == LANGUAGE_TOKEN_TYPE563    )564    is_boi_eoi[:, 0] |= tmp[:, 0] == VISION_TOKEN_TYPE565    is_boi_eoi[:, :-1] |= (tmp[:, :-1] == VISION_TOKEN_TYPE) & (566        tmp[:, 1:] == LANGUAGE_TOKEN_TYPE567    )568    is_boi_eoi[:, -1] |= tmp[:, -1] == VISION_TOKEN_TYPE569    tmp[is_boi_eoi] = LANGUAGE_TOKEN_TYPE570    # final position ids571    y = torch.zeros_like(x, dtype=torch.long)572    y[:, 1:] = (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE) | (573        (tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE)574    )575    y = y.cumsum(dim=-1)576    return y577 578 579class CogAgentModel(CogAgentPreTrainedModel):580    def __init__(self, config):581        super().__init__(config)582        self.padding_idx = config.pad_token_id583        self.vocab_size = config.vocab_size584 585        self.embed_tokens = nn.Embedding(586            config.vocab_size, config.hidden_size, self.padding_idx587        )588        self.layers = nn.ModuleList(589            [CogAgentDecoderLayer(config) for _ in range(config.num_hidden_layers)]590        )591        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)592 593        self.vision = EVA2CLIPModel(config)594        self.cross_vision = CrossVisionModel(config)595 596        self.gradient_checkpointing = False597        # Initialize weights and apply final processing598        self.post_init()599 600    def encode_images(self, images: List[List[torch.Tensor]]) -> torch.Tensor:601        images_list, images = images, []602 603        images = []604        for image_list in images_list:605            for image in image_list:606                images.append(image)607 608        images = torch.stack(images)609        images_features = self.vision(images)610        return images_features611 612    def encode_cross_images(self, images: List[List[torch.Tensor]]) -> torch.Tensor:613        images_list, images = images, []614 615        images = []616        for image_list in images_list:617            for image in image_list:618                images.append(image)619 620        images = torch.stack(images)621        encoder_outputs = self.cross_vision(images)622        return encoder_outputs623 624    def forward(625        self,626        input_ids: torch.LongTensor = None,627        images: List[List[torch.Tensor]] = None,628        cross_images: List[List[torch.Tensor]] = None,629        token_type_ids: Optional[torch.LongTensor] = None,630        attention_mask: Optional[torch.Tensor] = None,631        cross_attention_mask: Optional[torch.Tensor] = None,632        position_ids: Optional[torch.LongTensor] = None,633        past_key_values: Optional[List[torch.FloatTensor]] = None,634        inputs_embeds: Optional[torch.FloatTensor] = None,635        use_cache: Optional[bool] = None,636        output_attentions: Optional[bool] = None,637        output_hidden_states: Optional[bool] = None,638        return_dict: Optional[bool] = None,639    ) -> Union[Tuple, BaseModelOutputWithPast]:640        """take care of image_encode, token_type_ids, position_ids and (attention_mask = None is fine)"""641 642        if past_key_values is not None:643            encoder_outputs = None644            # generate mode with past_key_values. the image features are already mapped645        else:646            # not allow for inputs_embeds, because we want to process image feature647            assert (648                input_ids is not None and inputs_embeds is None649            ), f"{input_ids} {inputs_embeds}"650            if not is_empty(images):  # multi-modality651                assert (652                    token_type_ids is not None653                ), f"multi-modality requires `token_type_ids`!"654                assert len(input_ids) == len(images), f"{len(input_ids)} {len(images)}"655                inputs_embeds = self.embed_tokens(input_ids)656                images_features = self.encode_images(images)657                encoder_outputs = self.encode_cross_images(cross_images)658                images_features = rearrange(images_features, "b n d -> (b n) d")659                images_features = images_features.to(660                    dtype=inputs_embeds.dtype, device=inputs_embeds.device661                )662                inputs_embeds = inputs_embeds.index_put(663                    [token_type_ids == VISION_TOKEN_TYPE], images_features664                )665            else:  # single-modality666                if token_type_ids is None:667                    token_type_ids = (668                        torch.ones_like(669                            input_ids, dtype=torch.long, device=input_ids.device670                        )671                        * LANGUAGE_TOKEN_TYPE672                    )673                assert not (674                    token_type_ids == VISION_TOKEN_TYPE675                ).any(), f"{(token_type_ids == VISION_TOKEN_TYPE).sum()}"676                inputs_embeds = self.embed_tokens(input_ids)677                encoder_outputs = None678 679            if position_ids is None:680                position_ids = build_position_ids(token_type_ids, attention_mask)681            input_ids = None682 683        return self.llm_forward(684            input_ids=input_ids,685            encoder_outputs=encoder_outputs,686            token_type_ids=token_type_ids,687            attention_mask=attention_mask,688            cross_attention_mask=cross_attention_mask,689            position_ids=position_ids,690            past_key_values=past_key_values,691            inputs_embeds=inputs_embeds,692            use_cache=use_cache,693            output_attentions=output_attentions,694            output_hidden_states=output_hidden_states,695            return_dict=return_dict,696        )697 698    def llm_forward(699        self,700        input_ids: torch.LongTensor = None,701        encoder_outputs: torch.LongTensor = None,702        token_type_ids: torch.LongTensor = None,703        attention_mask: Optional[torch.Tensor] = None,704        cross_attention_mask: Optional[torch.Tensor] = None,705        position_ids: Optional[torch.LongTensor] = None,706        past_key_values: Optional[List[torch.FloatTensor]] = None,707        inputs_embeds: Optional[torch.FloatTensor] = None,708        use_cache: Optional[bool] = None,709        output_attentions: Optional[bool] = None,710        output_hidden_states: Optional[bool] = None,711        return_dict: Optional[bool] = None,712    ) -> Union[Tuple, BaseModelOutputWithPast]:713        """largely copy from llama forward and adapt for CogAgent with `token_type_ids`"""714        output_attentions = (715            output_attentions716            if output_attentions is not None717            else self.config.output_attentions718        )719        output_hidden_states = (720            output_hidden_states721            if output_hidden_states is not None722            else self.config.output_hidden_states723        )724        use_cache = use_cache if use_cache is not None else self.config.use_cache725 726        return_dict = (727            return_dict if return_dict is not None else self.config.use_return_dict728        )729 730        # retrieve input_ids and inputs_embeds731        if input_ids is not None and inputs_embeds is not None:732            raise ValueError(733                "You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time"734            )735        elif input_ids is not None:736            batch_size, seq_length = input_ids.shape737        elif inputs_embeds is not None:738            batch_size, seq_length, _ = inputs_embeds.shape739        else:740            raise ValueError(741                "You have to specify either decoder_input_ids or decoder_inputs_embeds"742            )743 744        seq_length_with_past = seq_length745        past_key_values_length = 0746 747        if past_key_values is not None:748            past_key_values_length = past_key_values[0][0].shape[2]749            seq_length_with_past = seq_length_with_past + past_key_values_length750 751        if position_ids is None:752            device = input_ids.device if input_ids is not None else inputs_embeds.device753            position_ids = torch.arange(754                past_key_values_length,755                seq_length + past_key_values_length,756                dtype=torch.long,757                device=device,758            )759            position_ids = position_ids.unsqueeze(0).view(-1, seq_length)760        else:761            position_ids = position_ids.view(-1, seq_length).long()762 763        if inputs_embeds is None:764            inputs_embeds = self.embed_tokens(input_ids)765        # embed positions766        if attention_mask is None:767            attention_mask = torch.ones(768                (batch_size, seq_length_with_past),769                dtype=torch.bool,770                device=inputs_embeds.device,771            )772        if cross_attention_mask is None:773            cross_attention_mask = torch.ones(774                (batch_size, 1), dtype=torch.bool, device=inputs_embeds.device775            )776        attention_mask = self._prepare_decoder_attention_mask(777            attention_mask,778            (batch_size, seq_length),779            inputs_embeds,780            past_key_values_length,781        )782 783        hidden_states = inputs_embeds784 785        # decoder layers786        all_hidden_states = () if output_hidden_states else None787        all_self_attns = () if output_attentions else None788        next_decoder_cache = () if use_cache else None789 790        for idx, decoder_layer in enumerate(self.layers):791            if output_hidden_states:792                all_hidden_states += (hidden_states,)793 794            past_key_value = (795                past_key_values[idx] if past_key_values is not None else None796            )797            layer_outputs = decoder_layer(798                hidden_states,799                encoder_outputs=encoder_outputs,800                token_type_ids=token_type_ids,801                attention_mask=attention_mask,802                cross_attention_mask=cross_attention_mask,803                position_ids=position_ids,804                past_key_value=past_key_value,805                output_attentions=output_attentions,806                use_cache=use_cache,807            )808            hidden_states = layer_outputs[0]809 810            if use_cache:811                next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)812 813            if output_attentions:814                all_self_attns += (layer_outputs[1],)815 816        hidden_states = self.norm(hidden_states)817 818        # add hidden states from the last decoder layer819        if output_hidden_states:820            all_hidden_states += (hidden_states,)821 822        next_cache = next_decoder_cache if use_cache else None823        if not return_dict:824            return tuple(825                v826                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]827                if v is not None828            )829        return BaseModelOutputWithPast(830            last_hidden_state=hidden_states,831            past_key_values=next_cache,832            hidden_states=all_hidden_states,833            attentions=all_self_attns,834        )835 836    def get_input_embeddings(self):837        return self.embed_tokens838 839    def set_input_embeddings(self, value):840        self.embed_tokens = value841 842    # noinspection PyMethodMayBeStatic843    # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask844    def _prepare_decoder_attention_mask(845        self, attention_mask, input_shape, inputs_embeds, past_key_values_length846    ):847        # create causal mask848        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]849        combined_attention_mask = None850        if input_shape[-1] > 1:851            combined_attention_mask = _make_causal_mask(852                input_shape,853                inputs_embeds.dtype,854                device=inputs_embeds.device,855                past_key_values_length=past_key_values_length,856            )857 858        if attention_mask is not None:859            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]860            expanded_attn_mask = _expand_mask(861                attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]862            ).to(inputs_embeds.device)863            combined_attention_mask = (864                expanded_attn_mask865                if combined_attention_mask is None866                else expanded_attn_mask + combined_attention_mask867            )868 869        return combined_attention_mask870 871 872def vqa_history_to_prompt(history, query):873    # Only support single round chat in vqa mode874    prompt = "<EOI>Question: "875    # for i, (old_query, response) in enumerate(history):876    #     prompt += old_query + " Short answer: " + response + " Question: "877    prompt += query + " Short answer:"878    return prompt879 880 881def chat_old_history_to_prompt(history, query):882    prompt = "<EOI>Question: "883    for i, (old_query, response) in enumerate(history):884        prompt += old_query + " Answer: " + response + "\nQuestion: "885    prompt += query + " Answer:"886    return prompt887 888 889def chat_history_to_prompt(history, query):890    prompt = " [INST] "891    for i, (old_query, response) in enumerate(history):892        prompt += old_query + " [/INST] " + response + " [INST] "893    prompt += query + " [/INST] "894    return prompt895 896 897def base_history_to_prompt(history, query):898    prompt = query899    return prompt900 901 902_history_to_prompt = {903    "base": base_history_to_prompt,904    "chat": chat_history_to_prompt,905    "chat_old": chat_old_history_to_prompt,906    "vqa": vqa_history_to_prompt,907}908 909 910class CogAgentForCausalLM(CogAgentPreTrainedModel):911    _auto_class = "AutoModelForCausalLM"912 913    def __init__(self, config):914        super().__init__(config)915        self.model = CogAgentModel(config)916        self.vocab_size = config.vocab_size917        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)918 919        # Initialize weights and apply final processing920        self.post_init()921 922    def get_input_embeddings(self):923        return self.model.embed_tokens924 925    def set_input_embeddings(self, value):926        self.model.embed_tokens = value927 928    def get_output_embeddings(self):929        return self.lm_head930 931    def set_output_embeddings(self, new_embeddings):932        self.lm_head = new_embeddings933 934    def set_decoder(self, decoder):935        self.model = decoder936 937    def get_decoder(self):938        return self.model939 940    def forward(941        self,942        input_ids: torch.LongTensor = None,943        images: List[List[torch.Tensor]] = None,944        cross_images: List[List[torch.Tensor]] = None,945        token_type_ids: Optional[torch.LongTensor] = None,946        attention_mask: Optional[torch.Tensor] = None,947        position_ids: Optional[torch.LongTensor] = None,948        past_key_values: Optional[List[torch.FloatTensor]] = None,949        inputs_embeds: Optional[torch.FloatTensor] = None,950        use_cache: Optional[bool] = None,951        output_attentions: Optional[bool] = None,952        output_hidden_states: Optional[bool] = None,953        return_dict: Optional[bool] = None,954        labels: Optional[torch.LongTensor] = None,955    ) -> Union[Tuple, CausalLMOutputWithPast]:956        output_attentions = (957            output_attentions958            if output_attentions is not None959            else self.config.output_attentions960        )961        output_hidden_states = (962            output_hidden_states963            if output_hidden_states is not None964            else self.config.output_hidden_states965        )966        return_dict = (967            return_dict if return_dict is not None else self.config.use_return_dict968        )969 970        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)971        outputs = self.model(972            input_ids=input_ids,973            images=images,974            cross_images=cross_images,975            token_type_ids=token_type_ids,976            attention_mask=attention_mask,977            position_ids=position_ids,978            past_key_values=past_key_values,979            inputs_embeds=inputs_embeds,980            use_cache=use_cache,981            output_attentions=output_attentions,982            output_hidden_states=output_hidden_states,983            return_dict=return_dict,984        )985 986        hidden_states = outputs[0]987        logits = self.lm_head(hidden_states)988        logits = logits.float()989 990        loss = None991        if labels is not None:992            # Shift so that tokens < n predict n993            shift_logits = logits[..., :-1, :].contiguous()994            shift_labels = labels[..., 1:].contiguous()995            # Flatten the tokens996            loss_fct = CrossEntropyLoss()997            shift_logits = shift_logits.view(-1, self.config.vocab_size)998            shift_labels = shift_labels.view(-1)999            # Enable model parallelism1000            shift_labels = shift_labels.to(shift_logits.device)1001            loss = loss_fct(shift_logits, shift_labels)1002 1003        if not return_dict:1004            output = (logits,) + outputs[1:]1005            return (loss,) + output if loss is not None else output1006 1007        return CausalLMOutputWithPast(1008            loss=loss,1009            logits=logits,1010            past_key_values=outputs.past_key_values,1011            hidden_states=outputs.hidden_states,1012            attentions=outputs.attentions,1013        )1014 1015    def _prepare_attention_mask_for_generation(1016        self,1017        inputs: torch.Tensor,1018        pad_token_id: Optional[int],1019        eos_token_id: Optional[Union[int, List[int]]],1020    ) -> torch.LongTensor:1021        return torch.ones(inputs.shape[:2], dtype=torch.long, device=inputs.device)  # type: ignore1022 1023    def prepare_inputs_for_generation(1024        self,1025        input_ids,1026        token_type_ids,1027        images=None,1028        cross_images=None,1029        past_key_values=None,1030        attention_mask=None,1031        inputs_embeds=None,1032        **kwargs,1033    ):1034        # build position_ids if needed1035        position_ids = kwargs.get("position_ids", None)1036        if position_ids is None:1037            position_ids = build_position_ids(token_type_ids, attention_mask)1038 1039        if past_key_values:1040            input_ids = input_ids[:, -1:]1041            token_type_ids = token_type_ids[:, -1:]1042            position_ids = position_ids[:, -1:]1043 1044        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step1045        if inputs_embeds is not None and past_key_values is None:1046            model_inputs = {"inputs_embeds": inputs_embeds}1047        else:1048            model_inputs = {"input_ids": input_ids}1049 1050        model_inputs.update(1051            {1052                "token_type_ids": token_type_ids,1053                "images": images,1054                "cross_images": cross_images,1055                "position_ids": position_ids,1056                "past_key_values": past_key_values,1057                "use_cache": kwargs.get("use_cache"),1058                "attention_mask": attention_mask,1059            }1060        )1061        return model_inputs1062 1063    def _update_model_kwargs_for_generation(1064        self,1065        outputs: "ModelOutput",1066        model_kwargs: Dict[str, Any],1067        is_encoder_decoder: bool = False,1068        standardize_cache_format: bool = False,1069    ) -> Dict[str, Any]:1070        # update past_key_values1071        model_kwargs["past_key_values"] = self._extract_past_from_model_output(1072            outputs, standardize_cache_format=standardize_cache_format1073        )1074        if getattr(outputs, "state", None) is not None:1075            model_kwargs["state"] = outputs.state1076 1077        # update token_type_ids with last value1078        if "token_type_ids" in model_kwargs:1079            token_type_ids = model_kwargs["token_type_ids"]1080            new_token_type_ids = (1081                torch.ones(1082                    size=(token_type_ids.shape[0], 1),1083                    dtype=token_type_ids.dtype,1084                    device=token_type_ids.device,1085                )1086                * LANGUAGE_TOKEN_TYPE1087            )1088            model_kwargs["token_type_ids"] = torch.cat(1089                [token_type_ids, new_token_type_ids], dim=-11090            )1091 1092        if not is_encoder_decoder:1093            # update attention mask1094            if "attention_mask" in model_kwargs:1095                attention_mask = model_kwargs["attention_mask"]1096                model_kwargs["attention_mask"] = torch.cat(1097                    [1098                        attention_mask,1099                        attention_mask.new_ones((attention_mask.shape[0], 1)),1100                    ],1101                    dim=-1,1102                )1103        else:1104            # update decoder attention mask1105            if "decoder_attention_mask" in model_kwargs:1106                decoder_attention_mask = model_kwargs["decoder_attention_mask"]1107                model_kwargs["decoder_attention_mask"] = torch.cat(1108                    [1109                        decoder_attention_mask,1110                        decoder_attention_mask.new_ones(1111                            (decoder_attention_mask.shape[0], 1)1112                        ),1113                    ],1114                    dim=-1,1115                )1116 1117        return model_kwargs1118 1119    def _reorder_cache(self, past_key_values, beam_idx):1120        reordered_past = ()1121        for layer_past in past_key_values:1122            reordered_past += (1123                tuple(1124                    past_state.index_select(0, beam_idx.to(past_state.device))1125                    for past_state in layer_past1126                ),1127            )1128        return reordered_past1129 1130    def build_conversation_input_ids(1131        self,1132        tokenizer: "PreTrainedTokenizer",1133        *,1134        query: str,1135        history: Optional[List[Tuple[str, str]]] = None,1136        images: Optional[List["PIL.Image"]] = None,1137        template_version: Optional[Literal["base", "chat", "vqa"]] = None,1138    ):1139        image_size: int = self.config.vision_config["image_size"]1140        cross_image_size: int = self.config.cross_image_size1141        patch_size: int = self.config.vision_config["patch_size"]1142        template_version = template_version or self.config.template_version1143        assert images is None or len(images) <= 1, f"not support multi images by now."1144        history = history or []1145        text = _history_to_prompt[template_version](history, query)1146 1147        input_ids = [tokenizer.bos_token_id]1148        token_type_ids = [LANGUAGE_TOKEN_TYPE]1149        if images is not None and len(images) == 1:1150            ori = images1151            # vision1152            transform = transforms.Compose(1153                [1154                    transforms.Resize(1155                        (image_size, image_size),1156                        interpolation=transforms.InterpolationMode.BICUBIC,1157                    ),1158                    transforms.ToTensor(),1159                    transforms.Normalize(1160                        (0.48145466, 0.4578275, 0.40821073),1161                        (0.26862954, 0.26130258, 0.27577711),1162                    ),1163                ]1164            )1165            images = [transform(ori[0])]1166            cross_transform = transforms.Compose(1167                [1168                    transforms.Resize(1169                        (cross_image_size, cross_image_size),1170                        interpolation=transforms.InterpolationMode.BICUBIC,1171                    ),1172                    transforms.ToTensor(),1173                    transforms.Normalize(1174                        (0.48145466, 0.4578275, 0.40821073),1175                        (0.26862954, 0.26130258, 0.27577711),1176                    ),1177                ]1178            )1179            cross_images = [cross_transform(ori[0])]1180            # language1181            vision_token_num = (image_size // patch_size) * (1182                image_size // patch_size1183            ) + 21184            input_ids += [tokenizer.pad_token_id] * vision_token_num1185            token_type_ids += [VISION_TOKEN_TYPE] * vision_token_num1186        text_ids = tokenizer.encode(text, add_special_tokens=False)1187 1188        input_ids += text_ids1189        token_type_ids += [LANGUAGE_TOKEN_TYPE] * len(text_ids)1190        attention_mask = [1] * len(input_ids)1191 1192        return {1193            "input_ids": torch.tensor(input_ids, dtype=torch.long),1194            "token_type_ids": torch.tensor(token_type_ids, dtype=torch.long),1195            "attention_mask": torch.tensor(attention_mask, dtype=torch.long),1196            "images": images,1197            "cross_images": cross_images,1198        }1199