HugoLaurencon
HF staff
Adapt inputs_merger() Function for Scenarios Without Image Input (#5)
990903f
verified
# coding=utf-8 | |
# Copyright 2023 Mistral AI and the HuggingFace Inc. team. All rights reserved. | |
# | |
# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX | |
# and OPT implementations in this library. It has been modified from its | |
# original forms to accommodate minor architectural differences compared | |
# to GPT-NeoX and OPT used by the Meta AI team that trained the model. | |
# | |
# Licensed under the Apache License, Version 2.0 (the "License"); | |
# you may not use this file except in compliance with the License. | |
# You may obtain a copy of the License at | |
# | |
# http://www.apache.org/licenses/LICENSE-2.0 | |
# | |
# Unless required by applicable law or agreed to in writing, software | |
# distributed under the License is distributed on an "AS IS" BASIS, | |
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
# See the License for the specific language governing permissions and | |
# limitations under the License. | |
""" PyTorch VMistral model.""" | |
from dataclasses import dataclass | |
import inspect | |
import math | |
import warnings | |
from typing import List, Optional, Tuple, Union | |
import torch | |
import torch.nn.functional as F | |
import torch.utils.checkpoint | |
from torch import nn | |
from torch.nn import CrossEntropyLoss | |
from transformers.activations import ACT2FN | |
from transformers.modeling_attn_mask_utils import _prepare_4d_causal_attention_mask | |
from transformers.utils import ( | |
add_start_docstrings, | |
add_start_docstrings_to_model_forward, | |
is_flash_attn_2_available, | |
replace_return_docstrings, | |
) | |
from einops import rearrange, repeat | |
from transformers import PreTrainedModel | |
from transformers.utils import logging | |
from transformers.modeling_outputs import ModelOutput | |
from .configuration_vmistral import VMistralConfig | |
from .vision import SiglipVisionModel | |
if is_flash_attn_2_available(): | |
from flash_attn import flash_attn_func, flash_attn_varlen_func | |
from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa | |
_flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters) | |
logger = logging.get_logger(__name__) | |
_CONFIG_FOR_DOC = "VMistralConfig" | |
VMistral_PRETRAINED_MODEL_ARCHIVE_LIST = [ | |
"HuggingFaceM4/VLM_WebSight_finetuned" | |
] | |
class VMistralBaseModelOutputWithPast(ModelOutput): | |
""" | |
Base class for VMistral model's outputs that may also contain a past key/values (to speed up sequential decoding). | |
Args: | |
last_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. | |
If `past_key_values` is used only the last hidden-state of the sequences of shape `(batch_size, 1, | |
hidden_size)` is output. | |
past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): | |
Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape | |
`(batch_size, num_heads, sequence_length, embed_size_per_head)`) and optionally if | |
`config.is_encoder_decoder=True` 2 additional tensors of shape `(batch_size, num_heads, | |
encoder_sequence_length, embed_size_per_head)`. | |
Contains pre-computed hidden-states (key and values in the self-attention blocks and optionally if | |
`config.is_encoder_decoder=True` in the cross-attention blocks) that can be used (see `past_key_values` | |
input) to speed up sequential decoding. | |
hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): | |
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. | |
attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): | |
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, | |
sequence_length)`. | |
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention | |
heads. | |
image_hidden_states (`tuple(torch.FloatTensor)`, *optional*): | |
Tuple of `torch.FloatTensor` (one for the output of the image embeddings, `(batch_size, num_images, | |
sequence_length, hidden_size)`. | |
image_hidden_states of the model produced by the vision encoder, and optionally by the perceiver | |
""" | |
last_hidden_state: torch.FloatTensor = None | |
past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None | |
hidden_states: Optional[Tuple[torch.FloatTensor]] = None | |
attentions: Optional[Tuple[torch.FloatTensor]] = None | |
image_hidden_states: Optional[Tuple[torch.FloatTensor]] = None | |
class VMistralCausalLMOutputWithPast(ModelOutput): | |
""" | |
Base class for VMistral causal language model (or autoregressive) outputs. | |
Args: | |
loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): | |
Language modeling loss (for next-token prediction). | |
logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): | |
Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). | |
past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): | |
Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape | |
`(batch_size, num_heads, sequence_length, embed_size_per_head)`) | |
Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see | |
`past_key_values` input) to speed up sequential decoding. | |
hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): | |
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. | |
attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): | |
Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, | |
sequence_length)`. | |
Attentions weights after the attention softmax, used to compute the weighted average in the self-attention | |
heads. | |
image_hidden_states (`tuple(torch.FloatTensor)`, *optional*): | |
Tuple of `torch.FloatTensor` (one for the output of the image embeddings, `(batch_size, num_images, | |
sequence_length, hidden_size)`. | |
image_hidden_states of the model produced by the vision encoder, and optionally by the perceiver | |
""" | |
loss: Optional[torch.FloatTensor] = None | |
logits: torch.FloatTensor = None | |
past_key_values: Optional[List[torch.FloatTensor]] = None | |
hidden_states: Optional[Tuple[torch.FloatTensor]] = None | |
attentions: Optional[Tuple[torch.FloatTensor]] = None | |
image_hidden_states: Optional[Tuple[torch.FloatTensor]] = None | |
def expand_inputs_for_generation( | |
input_ids, | |
expand_size=1, | |
is_encoder_decoder=False, | |
attention_mask=None, | |
encoder_outputs=None, | |
**model_kwargs, | |
): | |
expanded_return_idx = ( | |
torch.arange(input_ids.shape[0]).view(-1, 1).repeat(1, expand_size).view(-1).to(input_ids.device) | |
) | |
input_ids = input_ids.index_select(0, expanded_return_idx) | |
model_kwargs["pixel_values"] = model_kwargs.get("pixel_values", None) | |
model_kwargs["image_hidden_states"] = model_kwargs.get("image_hidden_states", None) | |
if "token_type_ids" in model_kwargs: | |
token_type_ids = model_kwargs["token_type_ids"] | |
model_kwargs["token_type_ids"] = token_type_ids.index_select(0, expanded_return_idx) | |
if attention_mask is not None: | |
model_kwargs["attention_mask"] = attention_mask.index_select(0, expanded_return_idx) | |
if model_kwargs["pixel_values"] is not None: | |
model_kwargs["pixel_values"] = model_kwargs["pixel_values"].index_select(0, expanded_return_idx) | |
elif model_kwargs["image_hidden_states"] is not None: | |
model_kwargs["image_hidden_states"] = model_kwargs["image_hidden_states"].index_select(0, expanded_return_idx) | |
return input_ids, model_kwargs | |
def update_model_kwargs_for_generation(outputs, model_kwargs): | |
# must have this key set to at least None | |
if "past_key_values" in outputs: | |
model_kwargs["past_key_values"] = outputs.past_key_values | |
else: | |
model_kwargs["past_key_values"] = None | |
# update token_type_ids with last value | |
if "token_type_ids" in model_kwargs: | |
token_type_ids = model_kwargs["token_type_ids"] | |
model_kwargs["token_type_ids"] = torch.cat([token_type_ids, token_type_ids[:, -1].unsqueeze(-1)], dim=-1) | |
# update attention masks | |
if "attention_mask" in model_kwargs: | |
attention_mask = model_kwargs["attention_mask"] | |
model_kwargs["attention_mask"] = torch.cat( | |
[attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1 | |
) | |
# Get the precomputed image_hidden_states | |
model_kwargs["image_hidden_states"] = outputs.image_hidden_states | |
return model_kwargs | |
def prepare_inputs_for_generation(input_ids, past_key_values=None, **kwargs): | |
token_type_ids = kwargs.get("token_type_ids", None) | |
# only last token for inputs_ids if past is defined in kwargs | |
if past_key_values: | |
input_ids = input_ids[:, -1].unsqueeze(-1) | |
if token_type_ids is not None: | |
token_type_ids = token_type_ids[:, -1].unsqueeze(-1) | |
attention_mask = kwargs.get("attention_mask", None) | |
position_ids = kwargs.get("position_ids", None) | |
if attention_mask is not None and position_ids is None: | |
# create position_ids on the fly for batch generation | |
position_ids = attention_mask.long().cumsum(-1) - 1 | |
position_ids.masked_fill_(attention_mask == 0, 1) | |
if past_key_values: | |
position_ids = position_ids[:, -1].unsqueeze(-1) | |
pixel_values = kwargs.get("pixel_values", None) | |
image_hidden_states = kwargs.get("image_hidden_states", None) | |
return { | |
"input_ids": input_ids, | |
"past_key_values": past_key_values, | |
"use_cache": kwargs.get("use_cache"), | |
"position_ids": position_ids, | |
"attention_mask": attention_mask, | |
"token_type_ids": token_type_ids, | |
"pixel_values": pixel_values, | |
"image_hidden_states": image_hidden_states, | |
} | |
def freeze_model(model, module_exceptions=[]): | |
mapping = { | |
"LayerNorm": nn.LayerNorm, | |
"Linear": nn.Linear, | |
"Embedding": nn.Embedding, | |
} | |
module_exceptions_mapped = [mapping[m] for m in module_exceptions] | |
for module in model.modules(): | |
if module_exceptions and any([isinstance(module, t) for t in module_exceptions_mapped]): | |
module.requires_grad_(True) # Explicitly setting it to true to avoid any mistakes | |
else: | |
module.requires_grad_(False) | |
return model | |
class DecoupledEmbedding(nn.Embedding): | |
# Derived from https://pytorch.org/docs/stable/_modules/torch/nn/modules/sparse.html#Embedding | |
""" | |
Implements a decoupling of parameters to allow freezing (or not) a subset of the embeddings. | |
In practise, the regular `weight` can be trained or frozen (i.e. `partially_freeze=True`), and if `num_additional_embeddings` > 0, then it will create `num_additional_embeddings` additional parameters that are always trained. | |
If `num_additional_embeddings=0`, then the module defaults back to the regular behavior of `nn.Embedding`. | |
""" | |
def __init__( | |
self, | |
num_embeddings, | |
num_additional_embeddings, | |
embedding_dim, | |
partially_freeze=False, | |
device=None, | |
dtype=None, | |
padding_idx=None, | |
**kwargs, | |
) -> None: | |
""" | |
num_additional_embeddings: int. Number of additional embeddings. Only useful when you `partially_freeze=True`. | |
partially_freeze: bool. If True, the regular `weight` will be frozen. `additional_weight` is never frozen. | |
Note: there are a lot of other parameters to initialize a standard `nn.Embedding` such as `padding_idx`, `max_norm` or `norm_type`. We are not supporting these. | |
""" | |
if padding_idx is not None and padding_idx > num_embeddings: | |
raise ValueError(f"padding_idx must be within num_embeddings. Got {padding_idx} and {num_embeddings}") | |
super().__init__( | |
num_embeddings=num_embeddings, | |
embedding_dim=embedding_dim, | |
device=device, | |
dtype=dtype, | |
padding_idx=padding_idx, | |
**kwargs, | |
) | |
self.num_embeddings = num_embeddings | |
self.padding_idx = padding_idx | |
self.num_additional_embeddings = num_additional_embeddings | |
self.partially_freeze = partially_freeze | |
if partially_freeze: | |
self.weight.requires_grad_(False) | |
if self.num_additional_embeddings > 0: | |
self.additional_embedding = nn.Embedding( | |
num_embeddings=self.num_additional_embeddings, | |
embedding_dim=embedding_dim, | |
device=device, | |
dtype=dtype, | |
) | |
def forward(self, input_ids): | |
""" | |
we have 2 embeddings, with different indices - one pretrained self.weight and another | |
self.additional_embedding.weight that is being trained. | |
in order to make a lookup of the input ids, we: | |
1. find out the indices of the entries belonging to the 2nd embedding | |
2. extract those values while subtracting the size of the first embedding (num_embeddings), | |
since the 2nd embedding starts from 0 and not num_embeddings | |
3. perform the 2nd embedding lookup | |
4. now we handle the 1st embedding, we overwrite indices belonging to the 2nd embedding with a padding index | |
5. perform the 1st embedding lookup | |
6. now we overwrite the values in the 1st embedding lookup with the values of the 2nd embedding lookup | |
note: for the 1st embedding lookup we could have looked up only the low indices and not do | |
the padding, but then we have to create a new tensor and populate it with 2 tensors that are | |
spread out across various indices - i.e. not a simple concat - I haven't benchmarked the | |
complex case if it's any faster, given that seqlens are usually relatively short it's | |
probably not faster or if faster not by much - but might be a good idea to measure. | |
""" | |
if self.num_additional_embeddings == 0: | |
return self.additional_embedding(input_ids) | |
# Clone so that we don't modify the original input_ids later on | |
input_ids = input_ids.clone() | |
additional_vocab_indices = torch.where(input_ids >= self.num_embeddings) | |
input_ids_additional_vocab = input_ids[additional_vocab_indices] | |
additional_embeddings = self.additional_embedding(input_ids_additional_vocab - self.num_embeddings) | |
# for successful lookup replace input_ids with 0, the results of these will be discarded anyway | |
input_ids[additional_vocab_indices] = 0 | |
full_vector = F.embedding(input_ids, self.weight) | |
# overwrite the records with high indices | |
full_vector[additional_vocab_indices] = additional_embeddings | |
return full_vector | |
def extra_repr(self) -> str: | |
return "num_embeddings={}, num_additional_embeddings={}, embedding_dim={}, partially_freeze={}".format( | |
self.num_embeddings, | |
self.num_additional_embeddings, | |
self.embedding_dim, | |
self.partially_freeze, | |
) | |
class DecoupledLinear(nn.Linear): | |
# Derived from https://pytorch.org/docs/stable/_modules/torch/nn/modules/linear.html#Linear | |
""" | |
Implements a decoupling of parameters to allow freezing (or not) a subset of the parameters. | |
In practise, the regular `weight` can be trained or frozen (i.e. `partially_freeze=True`), and if `out_additional_features` > 0, then it will create `out_additional_features * in_features` additional parameters that are always trained. | |
If `out_additional_features=0`, then the module defaults back to the regular behavior of `nn.Linear`. | |
""" | |
def __init__( | |
self, | |
in_features: int, | |
out_features: int, | |
out_additional_features: int = 0, | |
bias: bool = True, | |
partially_freeze: bool = True, | |
device=None, | |
dtype=None, | |
) -> None: | |
""" | |
out_additional_features: int. Number of additional trainable dimensions. Only makes sense when `partially_freeze=True`. | |
partially_freeze: bool. If True, the regular `weight` will be frozen and extra parameters (if any) will be trainable. If False, default to the regular behavior of nn.Linear. | |
""" | |
super().__init__(in_features, out_features, bias, device, dtype) | |
self.out_additional_features = out_additional_features | |
self.partially_freeze = partially_freeze | |
self.in_features = in_features | |
self.out_features = out_features | |
if partially_freeze: | |
self.weight.requires_grad_(False) | |
if bias: | |
self.bias.requires_grad_(False) | |
if out_additional_features > 0: | |
self.additional_fc = nn.Linear( | |
in_features=in_features, | |
out_features=out_additional_features, | |
bias=bias, | |
device=device, | |
dtype=dtype, | |
) | |
def forward(self, input: torch.Tensor) -> torch.Tensor: | |
output = F.linear(input, self.weight, self.bias) | |
if self.out_additional_features > 0: | |
additional_features = self.additional_fc(input) | |
output = torch.cat((output, additional_features), -1) | |
return output | |
def extra_repr(self) -> str: | |
"""Overwriting `nn.Linear.extra_repr` to include new parameters.""" | |
return "in_features={}, out_features={}, out_additional_features={}, bias={}, partially_freeze={}".format( | |
self.in_features, | |
self.out_features, | |
self.out_additional_features, | |
self.bias is not None, | |
self.partially_freeze, | |
) | |
class SwiGLU(nn.Module): | |
def __init__(self, embed_dim) -> None: | |
super().__init__() | |
self.fc1 = nn.Linear(embed_dim, embed_dim, bias=False) | |
self.fc2 = nn.Linear(embed_dim, embed_dim, bias=False) | |
def forward(self, x: torch.Tensor) -> torch.Tensor: | |
x_1 = self.fc1(x) | |
x_1 = torch.mul(x_1, torch.sigmoid(x_1)) | |
x_2 = self.fc2(x) | |
x = torch.mul(x_1, x_2) | |
return x | |
class ModalityProjection(nn.Module): | |
def __init__(self, embed_dim_in, embed_dim_out) -> None: | |
super().__init__() | |
self.fc1 = nn.Linear(embed_dim_in, embed_dim_out, bias=False) | |
self.act = SwiGLU(embed_dim_out) | |
self.fc2 = nn.Linear(embed_dim_out, embed_dim_out, bias=False) | |
def forward(self, x: torch.Tensor) -> torch.Tensor: | |
x = self.fc1(x) | |
x = self.act(x) | |
x = self.fc2(x) | |
return x | |
class PerceiverResampler(nn.Module): | |
def __init__( | |
self, embed_dim: int, depth: int, n_heads: int, head_dim: int, n_latents: int, qk_layer_norms: bool | |
) -> None: | |
""" | |
Instantiates a Perceiver Resampler that operates over a sequence of embeddings (say from a ResNet or ViT or | |
MAE) of a given dimension, performs `depth` blocks of cross-attention with a fixed `n_latents` inputs, then | |
returns a Tensor of shape [bsz, n_latents, embed_dim]. | |
:param embed_dim: Dimensionality of embeddings being fed to the Perceiver Resampler (also dimensionality of | |
latent embeddings *returned* by the Perceiver Resampler. Could be e.g., VIT embed_dim, ResNet | |
pool dim, and so on. | |
:param depth: Depth of the Perceiver Resampler (Transformer w/ cross attention). Should be shallow (< 3). | |
:param n_heads: Number of heads in each Transformer block (for multi-headed self-attention). | |
:param head_dim: Dimensionality of each head projection in the Transformer block. | |
:param n_latents: Number of latent embeddings to resample ("compress") the input sequence to (usually < 128). | |
""" | |
super().__init__() | |
self.embed_dim, self.n_heads, self.head_dim, self.n_latents = embed_dim, n_heads, head_dim, n_latents | |
self.qk_layer_norms = qk_layer_norms | |
# Create Latents for Perceiver | |
self.latents = nn.Parameter(torch.ones(self.n_latents, self.embed_dim)) | |
self.intermediate_dim = self.embed_dim * 4 | |
# Create Transformer Blocks | |
self.blocks = nn.ModuleList( | |
[ | |
nn.ModuleList( | |
[ | |
PerceiverAttention(self.embed_dim, self.n_heads, self.head_dim, self.qk_layer_norms), | |
MLP(self.embed_dim, self.intermediate_dim), | |
] | |
) | |
for _ in range(depth) | |
] | |
) | |
self.layer_norm = nn.LayerNorm(self.embed_dim) | |
def forward(self, context: torch.Tensor) -> torch.Tensor: | |
"""Resample arbitrary length context & *compress* down to self.n_latents latent embeddings""" | |
latents = repeat(self.latents, "seq embed -> bsz seq embed", bsz=context.shape[0]) | |
# Feed through Perceiver Attention blocks... | |
for attn, ff in self.blocks: | |
latents = attn(context, latents) + latents | |
latents = ff(latents) + latents | |
return self.layer_norm(latents) | |
class PerceiverAttention(nn.Module): | |
def __init__(self, embed_dim: int, n_heads: int, head_dim: int, qk_layer_norms: bool) -> None: | |
"""Perceiver Cross-Attention Module --> let long-form inputs be `context`, resampled embeddings be `latents`""" | |
super().__init__() | |
self.embed_dim, self.n_heads, self.head_dim = embed_dim, n_heads, head_dim | |
self.qk_layer_norms = qk_layer_norms | |
# Normalization & Scaling | |
self.context_layer_norm = nn.LayerNorm(self.embed_dim) | |
self.latents_layer_norm = nn.LayerNorm(self.embed_dim) | |
if self.qk_layer_norms: | |
self.q_layer_norm = nn.LayerNorm(self.head_dim) | |
self.k_layer_norm = nn.LayerNorm(self.head_dim) | |
self.qk_scale = self.head_dim**-0.5 | |
# Q, K, V Projection (no bias -- detail from Perceiver/Flamingo Papers). | |
self.q_proj = nn.Linear(self.embed_dim, self.n_heads * self.head_dim, bias=False) | |
self.k_proj = nn.Linear(self.embed_dim, self.n_heads * self.head_dim, bias=False) | |
self.v_proj = nn.Linear(self.embed_dim, self.n_heads * self.head_dim, bias=False) | |
self.output_proj = nn.Linear(self.n_heads * self.head_dim, self.embed_dim, bias=False) | |
def forward(self, context: torch.Tensor, latents: torch.Tensor) -> torch.Tensor: | |
""" | |
Runs Perceiver Self-Attention, with special (context, latents) appended along the `seq` dimension! | |
:param context: Tensor of shape [bsz, seq, embed_dim] representing long-form context to resample. | |
:param latents: Tensor of shape [bsz, n_latents, embed_dim] representing fixed length latents to compress to. | |
:return: Tensor of shape [bsz, n_latents, embed_dim] representing attention over latents w/ cross from context. | |
""" | |
context = self.context_layer_norm(context) | |
latents = self.latents_layer_norm(latents) | |
# Query, Key, Value Projections --> Note that in Flamingo, latents are *concatenated* with context prior to attn! | |
# Note: This results in queries w/ `seq = n_latents`, and keys, values with `seq = len(context) + n_latents` | |
q = self.q_proj(latents) | |
k = self.k_proj(torch.cat([context, latents], dim=-2)) | |
v = self.v_proj(torch.cat([context, latents], dim=-2)) | |
# Multiheaded Self-Attention w/ stable softmax (subtract per-row max -- `amax` -- before softmax call) | |
# =>> `attn` should be a 2D matrix of shape [n_latents x (context + n_latents)] | |
q, k, v = [rearrange(x, "bsz seq (heads embed) -> bsz heads seq embed", heads=self.n_heads) for x in (q, k, v)] | |
if self.qk_layer_norms: | |
q = self.q_layer_norm(q) | |
k = self.k_layer_norm(k) | |
scores = torch.einsum("... i d, ... j d -> ... i j", q * self.qk_scale, k) | |
stabilized_scores = scores - (scores.amax(dim=-1, keepdim=True).detach()) | |
attn = stabilized_scores.softmax(dim=-1) | |
# Attend & project back to output... | |
resampled = torch.einsum("... i j, ... j d -> ... i d", attn, v) | |
return self.output_proj( | |
rearrange(resampled, "bsz heads seq embed -> bsz seq (heads embed)", heads=self.n_heads) | |
) | |
class MLP(nn.Module): | |
def __init__(self, embed_dim, intermediate_size): | |
"""Simple MLP block with intermediate_size and embedding size""" | |
super().__init__() | |
self.embed_dim = embed_dim | |
self.ln = nn.LayerNorm(self.embed_dim) | |
self.fc = nn.Linear(self.embed_dim, intermediate_size, bias=False) | |
self.act = nn.ReLU() | |
self.c_proj = nn.Linear(intermediate_size, self.embed_dim, bias=False) | |
def forward(self, hidden_states: Optional[Tuple[torch.FloatTensor]]) -> torch.FloatTensor: | |
hidden_states = self.ln(hidden_states) | |
hidden_states = self.fc(hidden_states) | |
hidden_states = self.act(hidden_states) | |
hidden_states = self.c_proj(hidden_states) | |
return hidden_states | |
# Copied from transformers.models.llama.modeling_llama._get_unpad_data | |
def _get_unpad_data(attention_mask): | |
seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) | |
indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() | |
max_seqlen_in_batch = seqlens_in_batch.max().item() | |
cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)) | |
return ( | |
indices, | |
cu_seqlens, | |
max_seqlen_in_batch, | |
) | |
# Copied from transformers.models.llama.modeling_llama.LlamaRMSNorm with Llama->Mistral | |
class MistralRMSNorm(nn.Module): | |
def __init__(self, hidden_size, eps=1e-6): | |
""" | |
MistralRMSNorm is equivalent to T5LayerNorm | |
""" | |
super().__init__() | |
self.weight = nn.Parameter(torch.ones(hidden_size)) | |
self.variance_epsilon = eps | |
def forward(self, hidden_states): | |
input_dtype = hidden_states.dtype | |
hidden_states = hidden_states.to(torch.float32) | |
variance = hidden_states.pow(2).mean(-1, keepdim=True) | |
hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) | |
return self.weight * hidden_states.to(input_dtype) | |
# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Mistral | |
class MistralRotaryEmbedding(nn.Module): | |
def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): | |
super().__init__() | |
self.dim = dim | |
self.max_position_embeddings = max_position_embeddings | |
self.base = base | |
inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)) | |
self.register_buffer("inv_freq", inv_freq, persistent=False) | |
# Build here to make `torch.jit.trace` work. | |
self._set_cos_sin_cache( | |
seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype() | |
) | |
def _set_cos_sin_cache(self, seq_len, device, dtype): | |
self.max_seq_len_cached = seq_len | |
t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) | |
freqs = torch.einsum("i,j->ij", t, self.inv_freq) | |
# Different from paper, but it uses a different permutation in order to obtain the same calculation | |
emb = torch.cat((freqs, freqs), dim=-1) | |
self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False) | |
self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False) | |
def forward(self, x, seq_len=None): | |
# x: [bs, num_attention_heads, seq_len, head_size] | |
if seq_len > self.max_seq_len_cached: | |
self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) | |
return ( | |
self.cos_cached[:seq_len].to(dtype=x.dtype), | |
self.sin_cached[:seq_len].to(dtype=x.dtype), | |
) | |
# Copied from transformers.models.llama.modeling_llama.rotate_half | |
def rotate_half(x): | |
"""Rotates half the hidden dims of the input.""" | |
x1 = x[..., : x.shape[-1] // 2] | |
x2 = x[..., x.shape[-1] // 2 :] | |
return torch.cat((-x2, x1), dim=-1) | |
# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb | |
def apply_rotary_pos_emb(q, k, cos, sin, position_ids): | |
cos = cos[position_ids].unsqueeze(1) # [seq_len, dim] -> [batch_size, 1, seq_len, head_dim] | |
sin = sin[position_ids].unsqueeze(1) | |
q_embed = (q * cos) + (rotate_half(q) * sin) | |
k_embed = (k * cos) + (rotate_half(k) * sin) | |
return q_embed, k_embed | |
class MistralMLP(nn.Module): | |
def __init__(self, config): | |
super().__init__() | |
self.config = config | |
self.hidden_size = config.hidden_size | |
self.intermediate_size = config.intermediate_size | |
self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) | |
self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) | |
self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) | |
self.act_fn = ACT2FN[config.hidden_act] | |
def forward(self, x): | |
return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) | |
# Copied from transformers.models.llama.modeling_llama.repeat_kv | |
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: | |
""" | |
This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, | |
num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) | |
""" | |
batch, num_key_value_heads, slen, head_dim = hidden_states.shape | |
if n_rep == 1: | |
return hidden_states | |
hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) | |
return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim) | |
class MistralAttention(nn.Module): | |
""" | |
Multi-headed attention from 'Attention Is All You Need' paper. Modified to use sliding window attention: Longformer | |
and "Generating Long Sequences with Sparse Transformers". | |
""" | |
def __init__(self, config: VMistralConfig, qk_layer_norms: bool = False): | |
super().__init__() | |
self.config = config | |
self.hidden_size = config.hidden_size | |
self.num_heads = config.num_attention_heads | |
self.head_dim = self.hidden_size // self.num_heads | |
self.num_key_value_heads = config.num_key_value_heads | |
self.num_key_value_groups = self.num_heads // self.num_key_value_heads | |
self.max_position_embeddings = config.max_position_embeddings | |
self.rope_theta = config.rope_theta | |
self.is_causal = True | |
if (self.head_dim * self.num_heads) != self.hidden_size: | |
raise ValueError( | |
f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}" | |
f" and `num_heads`: {self.num_heads})." | |
) | |
self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False) | |
self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) | |
self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) | |
self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) | |
self.qk_layer_norms = qk_layer_norms | |
if self.qk_layer_norms: | |
self.q_layer_norm = MistralRMSNorm(self.head_dim, eps=config.rms_norm_eps) | |
self.k_layer_norm = MistralRMSNorm(self.head_dim, eps=config.rms_norm_eps) | |
self.rotary_emb = MistralRotaryEmbedding( | |
self.head_dim, | |
max_position_embeddings=self.max_position_embeddings, | |
base=self.rope_theta, | |
) | |
self.attention_dropout = config.attention_dropout | |
def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): | |
return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() | |
def forward( | |
self, | |
hidden_states: torch.Tensor, | |
key_value_states: Optional[torch.Tensor] = None, | |
attention_mask: Optional[torch.Tensor] = None, | |
position_ids: Optional[torch.LongTensor] = None, | |
past_key_value: Optional[Tuple[torch.Tensor]] = None, | |
output_attentions: bool = False, | |
use_cache: bool = False, | |
**kwargs, | |
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: | |
if "padding_mask" in kwargs: | |
warnings.warn( | |
"Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use" | |
" `attention_mask` instead.`" | |
) | |
bsz, q_len, _ = hidden_states.size() | |
query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) | |
key_states = ( | |
self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) | |
) | |
value_states = ( | |
self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) | |
) | |
kv_seq_len = key_states.shape[-2] | |
if past_key_value is not None: | |
kv_seq_len += past_key_value[0].shape[-2] | |
cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) | |
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) | |
if past_key_value is not None: | |
# reuse k, v, self_attention | |
key_states = torch.cat([past_key_value[0], key_states], dim=2) | |
value_states = torch.cat([past_key_value[1], value_states], dim=2) | |
past_key_value = (key_states, value_states) if use_cache else None | |
if self.qk_layer_norms: | |
query_states = self.q_layer_norm(query_states) | |
key_states = self.k_layer_norm(key_states) | |
# repeat k/v heads if n_kv_heads < n_heads | |
key_states = repeat_kv(key_states, self.num_key_value_groups) | |
value_states = repeat_kv(value_states, self.num_key_value_groups) | |
attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) | |
if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len): | |
raise ValueError( | |
f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is" | |
f" {attn_weights.size()}" | |
) | |
if attention_mask is not None: | |
if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): | |
raise ValueError( | |
f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}" | |
) | |
attn_weights = attn_weights + attention_mask | |
# upcast attention to fp32 | |
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) | |
attn_output = torch.matmul(attn_weights, value_states) | |
if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim): | |
raise ValueError( | |
f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is" | |
f" {attn_output.size()}" | |
) | |
attn_output = attn_output.transpose(1, 2).contiguous() | |
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size) | |
attn_output = self.o_proj(attn_output) | |
if not output_attentions: | |
attn_weights = None | |
return attn_output, attn_weights, past_key_value | |
class MistralFlashAttention2(MistralAttention): | |
""" | |
Mistral flash attention module. This module inherits from `MistralAttention` as the weights of the module stays | |
untouched. The only required change would be on the forward pass where it needs to correctly call the public API of | |
flash attention and deal with padding tokens in case the input contains any of them. | |
""" | |
def forward( | |
self, | |
hidden_states: torch.Tensor, | |
attention_mask: Optional[torch.Tensor] = None, | |
position_ids: Optional[torch.LongTensor] = None, | |
past_key_value: Optional[Tuple[torch.Tensor]] = None, | |
output_attentions: bool = False, | |
use_cache: bool = False, | |
**kwargs, | |
): | |
if "padding_mask" in kwargs: | |
warnings.warn( | |
"Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use" | |
" `attention_mask` instead.`" | |
) | |
# overwrite attention_mask with padding_mask | |
attention_mask = kwargs.pop("padding_mask") | |
bsz, q_len, _ = hidden_states.size() | |
query_states = self.q_proj(hidden_states) | |
key_states = self.k_proj(hidden_states) | |
value_states = self.v_proj(hidden_states) | |
query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) | |
key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) | |
value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) | |
kv_seq_len = key_states.shape[-2] | |
if past_key_value is not None: | |
kv_seq_len += past_key_value[0].shape[-2] | |
# Because the input can be padded, the absolute sequence length depends on the max position id. | |
rotary_seq_len = max(kv_seq_len, position_ids[:, -1].max().item()) + 1 | |
cos, sin = self.rotary_emb(value_states, seq_len=rotary_seq_len) | |
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) | |
use_sliding_windows = ( | |
_flash_supports_window_size | |
and hasattr(self.config, "sliding_window") is not None | |
and kv_seq_len > self.config.sliding_window | |
) | |
if not _flash_supports_window_size: | |
logger.warning_once( | |
"The current flash attention version does not support sliding window attention, for a more memory" | |
" efficient implementation make sure to upgrade flash-attn library." | |
) | |
if past_key_value is not None: | |
# Activate slicing cache only if the config has a value `sliding_windows` attribute | |
if hasattr(self.config, "sliding_window") and kv_seq_len > self.config.sliding_window: | |
slicing_tokens = kv_seq_len - self.config.sliding_window | |
past_key = past_key_value[0] | |
past_value = past_key_value[1] | |
past_key = past_key[:, :, slicing_tokens:, :].contiguous() | |
past_value = past_value[:, :, slicing_tokens:, :].contiguous() | |
if past_key.shape[-2] != self.config.sliding_window - 1: | |
raise ValueError( | |
"past key must have a shape of (`batch_size, num_heads, self.config.sliding_window-1," | |
f" head_dim`), got {past_key.shape}" | |
) | |
past_key_value = (past_key, past_value) | |
if attention_mask is not None: | |
attention_mask = attention_mask[:, slicing_tokens:] | |
attention_mask = torch.cat([attention_mask, torch.ones_like(attention_mask[:, -1:])], dim=-1) | |
key_states = torch.cat([past_key_value[0], key_states], dim=2) | |
value_states = torch.cat([past_key_value[1], value_states], dim=2) | |
past_key_value = (key_states, value_states) if use_cache else None | |
# repeat k/v heads if n_kv_heads < n_heads | |
key_states = repeat_kv(key_states, self.num_key_value_groups) | |
value_states = repeat_kv(value_states, self.num_key_value_groups) | |
dropout_rate = 0.0 if not self.training else self.attention_dropout | |
# In PEFT, usually we cast the layer norms in float32 for training stability reasons | |
# therefore the input hidden states gets silently casted in float32. Hence, we need | |
# cast them back in float16 just to be sure everything works as expected. | |
input_dtype = query_states.dtype | |
if input_dtype == torch.float32: | |
# Handle the case where the model is quantized | |
if hasattr(self.config, "_pre_quantization_dtype"): | |
target_dtype = self.config._pre_quantization_dtype | |
else: | |
target_dtype = self.q_proj.weight.dtype | |
logger.warning_once( | |
"The input hidden states seems to be silently casted in float32, this might be related to the fact" | |
" you have upcasted embedding or layer norm layers in float32. We will cast back the input in" | |
f" {target_dtype}." | |
) | |
query_states = query_states.to(target_dtype) | |
key_states = key_states.to(target_dtype) | |
value_states = value_states.to(target_dtype) | |
# Reashape to the expected shape for Flash Attention | |
query_states = query_states.transpose(1, 2) | |
key_states = key_states.transpose(1, 2) | |
value_states = value_states.transpose(1, 2) | |
attn_output = self._flash_attention_forward( | |
query_states, | |
key_states, | |
value_states, | |
attention_mask, | |
q_len, | |
dropout=dropout_rate, | |
use_sliding_windows=use_sliding_windows, | |
) | |
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() | |
attn_output = self.o_proj(attn_output) | |
if not output_attentions: | |
attn_weights = None | |
return attn_output, attn_weights, past_key_value | |
def _flash_attention_forward( | |
self, | |
query_states, | |
key_states, | |
value_states, | |
attention_mask, | |
query_length, | |
dropout=0.0, | |
softmax_scale=None, | |
use_sliding_windows=False, | |
): | |
""" | |
Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token | |
first unpad the input, then computes the attention scores and pad the final attention scores. | |
Args: | |
query_states (`torch.Tensor`): | |
Input query states to be passed to Flash Attention API | |
key_states (`torch.Tensor`): | |
Input key states to be passed to Flash Attention API | |
value_states (`torch.Tensor`): | |
Input value states to be passed to Flash Attention API | |
attention_mask (`torch.Tensor`): | |
The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the | |
position of padding tokens and 1 for the position of non-padding tokens. | |
dropout (`int`, *optional*): | |
Attention dropout | |
softmax_scale (`float`, *optional*): | |
The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim) | |
use_sliding_windows (`bool`, *optional*): | |
Whether to activate sliding window attention. | |
""" | |
# Contains at least one padding token in the sequence | |
if attention_mask is not None: | |
batch_size = query_states.shape[0] | |
query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input( | |
query_states, key_states, value_states, attention_mask, query_length | |
) | |
cu_seqlens_q, cu_seqlens_k = cu_seq_lens | |
max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens | |
if not use_sliding_windows: | |
attn_output_unpad = flash_attn_varlen_func( | |
query_states, | |
key_states, | |
value_states, | |
cu_seqlens_q=cu_seqlens_q, | |
cu_seqlens_k=cu_seqlens_k, | |
max_seqlen_q=max_seqlen_in_batch_q, | |
max_seqlen_k=max_seqlen_in_batch_k, | |
dropout_p=dropout, | |
softmax_scale=softmax_scale, | |
causal=self.is_causal, | |
) | |
else: | |
attn_output_unpad = flash_attn_varlen_func( | |
query_states, | |
key_states, | |
value_states, | |
cu_seqlens_q=cu_seqlens_q, | |
cu_seqlens_k=cu_seqlens_k, | |
max_seqlen_q=max_seqlen_in_batch_q, | |
max_seqlen_k=max_seqlen_in_batch_k, | |
dropout_p=dropout, | |
softmax_scale=softmax_scale, | |
causal=self.is_causal, | |
window_size=(self.config.sliding_window, self.config.sliding_window), | |
) | |
attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length) | |
else: | |
if not use_sliding_windows: | |
attn_output = flash_attn_func( | |
query_states, | |
key_states, | |
value_states, | |
dropout, | |
softmax_scale=softmax_scale, | |
causal=self.is_causal, | |
) | |
else: | |
attn_output = flash_attn_func( | |
query_states, | |
key_states, | |
value_states, | |
dropout, | |
softmax_scale=softmax_scale, | |
causal=self.is_causal, | |
window_size=(self.config.sliding_window, self.config.sliding_window), | |
) | |
return attn_output | |
def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length): | |
batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape | |
# On the first iteration we need to properly re-create the padding mask | |
# by slicing it on the proper place | |
if kv_seq_len != attention_mask.shape[-1]: | |
attention_mask_num_tokens = attention_mask.shape[-1] | |
attention_mask = attention_mask[:, attention_mask_num_tokens - kv_seq_len :] | |
indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask) | |
key_layer = index_first_axis(key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k) | |
value_layer = index_first_axis(value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k) | |
if query_length == kv_seq_len: | |
query_layer = index_first_axis( | |
query_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k | |
) | |
cu_seqlens_q = cu_seqlens_k | |
max_seqlen_in_batch_q = max_seqlen_in_batch_k | |
indices_q = indices_k | |
elif query_length == 1: | |
max_seqlen_in_batch_q = 1 | |
cu_seqlens_q = torch.arange( | |
batch_size + 1, dtype=torch.int32, device=query_layer.device | |
) # There is a memcpy here, that is very bad. | |
indices_q = cu_seqlens_q[:-1] | |
query_layer = query_layer.squeeze(1) | |
else: | |
# The -q_len: slice assumes left padding. | |
attention_mask = attention_mask[:, -query_length:] | |
query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, attention_mask) | |
return ( | |
query_layer, | |
key_layer, | |
value_layer, | |
indices_q, | |
(cu_seqlens_q, cu_seqlens_k), | |
(max_seqlen_in_batch_q, max_seqlen_in_batch_k), | |
) | |
class MistralDecoderLayer(nn.Module): | |
def __init__(self, config: VMistralConfig): | |
super().__init__() | |
self.hidden_size = config.hidden_size | |
self.self_attn = ( | |
MistralAttention(config=config) | |
if not getattr(config, "_flash_attn_2_enabled", False) | |
else MistralFlashAttention2(config) | |
) | |
self.mlp = MistralMLP(config) | |
self.input_layernorm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps) | |
self.post_attention_layernorm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps) | |
def forward( | |
self, | |
hidden_states: torch.Tensor, | |
attention_mask: Optional[torch.Tensor] = None, | |
position_ids: Optional[torch.LongTensor] = None, | |
past_key_value: Optional[Tuple[torch.Tensor]] = None, | |
output_attentions: Optional[bool] = False, | |
use_cache: Optional[bool] = False, | |
**kwargs, | |
) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: | |
if "padding_mask" in kwargs: | |
warnings.warn( | |
"Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use" | |
" `attention_mask` instead.`" | |
) | |
""" | |
Args: | |
hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` | |
attention_mask (`torch.FloatTensor`, *optional*): attention mask of size | |
`(batch, sequence_length)` where padding elements are indicated by 0. | |
output_attentions (`bool`, *optional*): | |
Whether or not to return the attentions tensors of all attention layers. See `attentions` under | |
returned tensors for more detail. | |
use_cache (`bool`, *optional*): | |
If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding | |
(see `past_key_values`). | |
past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states | |
""" | |
residual = hidden_states | |
hidden_states = self.input_layernorm(hidden_states) | |
# Self Attention | |
hidden_states, self_attn_weights, present_key_value = self.self_attn( | |
hidden_states=hidden_states, | |
attention_mask=attention_mask, | |
position_ids=position_ids, | |
past_key_value=past_key_value, | |
output_attentions=output_attentions, | |
use_cache=use_cache, | |
) | |
hidden_states = residual + hidden_states | |
# Fully Connected | |
residual = hidden_states | |
hidden_states = self.post_attention_layernorm(hidden_states) | |
hidden_states = self.mlp(hidden_states) | |
hidden_states = residual + hidden_states | |
outputs = (hidden_states,) | |
if output_attentions: | |
outputs += (self_attn_weights,) | |
if use_cache: | |
outputs += (present_key_value,) | |
return outputs | |
MISTRAL_START_DOCSTRING = r""" | |
This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the | |
library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads | |
etc.) | |
This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. | |
Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage | |
and behavior. | |
Parameters: | |
config ([`VMistralConfig`]): | |
Model configuration class with all the parameters of the model. Initializing with a config file does not | |
load the weights associated with the model, only the configuration. Check out the | |
[`~PreTrainedModel.from_pretrained`] method to load the model weights. | |
""" | |
class VMistralPreTrainedModel(PreTrainedModel): | |
config_class = VMistralConfig | |
base_model_prefix = "model" | |
supports_gradient_checkpointing = True | |
_no_split_modules = ["MistralDecoderLayer"] | |
_skip_keys_device_placement = "past_key_values" | |
_supports_sdpa = False | |
def _init_weights(self, module): | |
# important: this ported version of the model isn't meant for training from scratch - only | |
# inference and fine-tuning - so the proper init weights code has been removed - the m4 code | |
# base should be used for training from scratch and it contains the correct code. | |
std = self.config.initializer_range | |
if isinstance(module, nn.Linear): | |
module.weight.data.normal_(mean=0.0, std=std) | |
if module.bias is not None: | |
module.bias.data.zero_() | |
elif isinstance(module, nn.Embedding): | |
module.weight.data.normal_(mean=0.0, std=std) | |
if module.padding_idx is not None: | |
module.weight.data[module.padding_idx].zero_() | |
# @classmethod | |
# def override_vision_model_wrapper(cls, model, config, vision_model_name, vision_model_params, torch_dtype): | |
# # this can be called via from_pretrained from a class w/ head or w/o head so we extract the beheaded model version | |
# beheaded_model = model.model if hasattr(model, "model") else model | |
# cls.override_vision_model(beheaded_model, vision_model_name, vision_model_params, torch_dtype) | |
# beheaded_model.freeze_relevant_params(config) | |
MISTRAL_INPUTS_DOCSTRING = r""" | |
Args: | |
input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): | |
Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide | |
it. | |
Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and | |
[`PreTrainedTokenizer.__call__`] for details. | |
[What are input IDs?](../glossary#input-ids) | |
attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): | |
Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: | |
- 1 for tokens that are **not masked**, | |
- 0 for tokens that are **masked**. | |
[What are attention masks?](../glossary#attention-mask) | |
Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and | |
[`PreTrainedTokenizer.__call__`] for details. | |
If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see | |
`past_key_values`). | |
If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`] | |
and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more | |
information on the default strategy. | |
- 1 indicates the head is **not masked**, | |
- 0 indicates the head is **masked**. | |
position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): | |
Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0, | |
config.n_positions - 1]`. | |
[What are position IDs?](../glossary#position-ids) | |
past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): | |
Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape | |
`(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape | |
`(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. | |
Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention | |
blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. | |
If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that | |
don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all | |
`decoder_input_ids` of shape `(batch_size, sequence_length)`. | |
inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): | |
Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This | |
is useful if you want more control over how to convert `input_ids` indices into associated vectors than the | |
model's internal embedding lookup matrix. | |
use_cache (`bool`, *optional*): | |
If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see | |
`past_key_values`). | |
output_attentions (`bool`, *optional*): | |
Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned | |
tensors for more detail. | |
output_hidden_states (`bool`, *optional*): | |
Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for | |
more detail. | |
return_dict (`bool`, *optional*): | |
Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. | |
""" | |
class VMistralModel(VMistralPreTrainedModel): | |
""" | |
Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MistralDecoderLayer`] | |
Args: | |
config: VMistralConfig | |
""" | |
def __init__(self, config: VMistralConfig, vision_model=None): | |
super().__init__(config) | |
self.config = config | |
self.padding_idx = config.pad_token_id | |
self.vocab_size = config.vocab_size | |
self.sliding_window = config.sliding_window | |
self.embed_tokens = DecoupledEmbedding( | |
num_embeddings=config.vocab_size, | |
num_additional_embeddings=config.additional_vocab_size, | |
embedding_dim=config.hidden_size, | |
partially_freeze=config.freeze_text_layers, | |
padding_idx=self.padding_idx, | |
) | |
# Load an uninitialized model and later in from_pretrained will load the pre-trained model - | |
# this solves the losing of weights in `from_pretrained` on the main model | |
self.vision_model = SiglipVisionModel(config.vision_config) | |
# Dim projection - projecting from the vision dim to the text dim | |
self.modality_projection = ModalityProjection( | |
embed_dim_in=self.config.vision_config.hidden_size, embed_dim_out=self.config.hidden_size | |
) | |
# Perceiver Resampler | |
if config.use_resampler: | |
self.perceiver_resampler = PerceiverResampler( | |
config.hidden_size, | |
config.perceiver_config.resampler_depth, | |
config.perceiver_config.resampler_n_heads, | |
config.perceiver_config.resampler_head_dim, | |
config.perceiver_config.resampler_n_latents, | |
config.perceiver_config.qk_layer_norms_perceiver, | |
) | |
if config.use_resampler: | |
self.image_seq_len = config.perceiver_config.resampler_n_latents | |
else: | |
self.image_seq_len = ( | |
config.vision_config.image_size // config.vision_config.patch_size | |
) ** 2 # TODO: pretty sure that does not work for CLIP models since there is the CLS token | |
self.image_token_id = self.config.image_token_id | |
self.layers = nn.ModuleList([MistralDecoderLayer(config) for _ in range(config.num_hidden_layers)]) | |
self.gradient_checkpointing = False | |
self.norm = MistralRMSNorm(config.hidden_size, eps=config.rms_norm_eps) | |
# Initialize weights and apply final processing | |
self.post_init() | |
self.freeze_relevant_params(config) | |
def freeze_relevant_params(self, config=None): | |
if config is None: | |
config = self.config | |
if config.freeze_text_layers: | |
self.freeze_text_layers(config.freeze_text_module_exceptions) | |
if config.freeze_vision_layers: | |
freeze_model(self.vision_model, module_exceptions=config.freeze_vision_module_exceptions) | |
def freeze_text_layers(self, module_exceptions): | |
for module in [self.layers, self.norm]: | |
freeze_model(module, module_exceptions=module_exceptions) | |
def get_input_embeddings(self): | |
return self.embed_tokens | |
def set_input_embeddings(self, value): | |
self.embed_tokens = value | |
def inputs_merger( | |
self, | |
input_ids: torch.LongTensor = None, | |
inputs_embeds: Optional[torch.Tensor] = None, | |
image_hidden_states: Optional[torch.Tensor] = None, | |
): | |
""" | |
This method aims at merging the token embeddings with the image hidden states into one single sequence of vectors that are fed to the transformer LM. | |
The merging happens as follows: | |
- The text token sequence is: `tok_1 tok_2 tok_3 <fake_token_around_image> <image> <image> ... <image> <fake_token_around_image> tok_4`. | |
- We get the image hidden states for the image through the vision encoder (and potentially the perceiver), and that hidden state is then projected into the text embedding space. | |
We thus have a sequence of image hidden states of size (1, image_seq_len, hidden_dim), where 1 is for batch_size of 1 image and hidden_dim is the hidden_dim of the LM transformer. | |
- The merging happens so that we obtain the following sequence: `vector_tok_1 vector_tok_2 vector_tok_3 vector_fake_tok_around_image {sequence of image_seq_len image hidden states} vector_fake_toke_around_image vector_tok_4`. That sequence is fed to the LM. | |
- To fit the format of that sequence, `input_ids`, `input_embeds`, `attention_mask` are all 3 adapted to insert the image hidden states. | |
""" | |
batch_size = input_ids.size(0) | |
if inputs_embeds is not None: | |
new_inputs_embeds = inputs_embeds.clone() | |
if image_hidden_states is not None: | |
vision_pipeline_output_seq_len = image_hidden_states.shape[1] | |
vision_hidden_size = image_hidden_states.shape[2] | |
# Get the number of images for each example | |
num_images = (input_ids == self.image_token_id).sum(dim=-1) // self.image_seq_len | |
cum_num_images = num_images.cumsum(dim=-1) | |
for batch_idx in range(batch_size): | |
# Get the number of images for this particular example | |
example_num_images = num_images[batch_idx] | |
# Get the image_hidden_states corresponding to True images for the example, so get rid of the padding images. | |
start = 0 if batch_idx == 0 else cum_num_images[batch_idx - 1] | |
end = cum_num_images[batch_idx] | |
example_true_image_hidden_states = image_hidden_states[start:end] | |
if ( | |
new_inputs_embeds[batch_idx][input_ids[batch_idx] == self.image_token_id].shape[0] | |
!= example_num_images * vision_pipeline_output_seq_len | |
): | |
raise ValueError( | |
"new_inputs_embeds to replace has shape[0]:" | |
f" {new_inputs_embeds[batch_idx][input_ids[batch_idx] == self.image_token_id].shape[0]} but" | |
" should have shape[0]:" | |
f" {example_num_images}*{vision_pipeline_output_seq_len}={example_num_images * vision_pipeline_output_seq_len} " | |
) | |
# Insert the image_hidden_states | |
new_inputs_embeds[batch_idx][input_ids[batch_idx] == self.image_token_id] = ( | |
example_true_image_hidden_states.view( | |
example_num_images * vision_pipeline_output_seq_len, | |
vision_hidden_size, | |
) | |
) | |
return_dict = {} | |
if inputs_embeds is not None: | |
return_dict["inputs_embeds"] = new_inputs_embeds | |
return return_dict | |
def forward( | |
self, | |
input_ids: torch.LongTensor = None, | |
attention_mask: Optional[torch.Tensor] = None, | |
position_ids: Optional[torch.LongTensor] = None, | |
past_key_values: Optional[List[torch.FloatTensor]] = None, | |
inputs_embeds: Optional[torch.FloatTensor] = None, | |
pixel_values: Optional[torch.FloatTensor] = None, | |
image_hidden_states: Optional[torch.FloatTensor] = None, | |
use_cache: Optional[bool] = None, | |
output_attentions: Optional[bool] = None, | |
output_hidden_states: Optional[bool] = None, | |
return_dict: Optional[bool] = None, | |
) -> Union[Tuple, VMistralBaseModelOutputWithPast]: | |
device = input_ids.device if input_ids is not None else inputs_embeds.device | |
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions | |
output_hidden_states = ( | |
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states | |
) | |
use_cache = use_cache if use_cache is not None else self.config.use_cache | |
return_dict = return_dict if return_dict is not None else self.config.use_return_dict | |
# retrieve input_ids and inputs_embeds | |
if input_ids is not None and inputs_embeds is not None: | |
raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") | |
elif input_ids is not None: | |
batch_size, seq_length = input_ids.shape | |
elif inputs_embeds is not None: | |
batch_size, seq_length, _ = inputs_embeds.shape | |
else: | |
raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") | |
seq_length_with_past = seq_length | |
past_key_values_length = 0 | |
if past_key_values is not None: | |
past_key_values_length = past_key_values[0][0].shape[2] | |
seq_length_with_past = seq_length_with_past + past_key_values_length | |
if position_ids is None: | |
device = input_ids.device if input_ids is not None else inputs_embeds.device | |
position_ids = torch.arange( | |
past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device | |
) | |
position_ids = position_ids.unsqueeze(0).view(-1, seq_length) | |
else: | |
position_ids = position_ids.view(-1, seq_length).long() | |
if inputs_embeds is None: | |
inputs_embeds = self.embed_tokens(input_ids) | |
# START VISUAL INPUTS INTEGRATION | |
if pixel_values is not None and image_hidden_states is not None: | |
raise ValueError("You cannot specify both pixel_values and image_hidden_states at the same time") | |
elif pixel_values is not None: | |
pixel_values = pixel_values.to(dtype=self.dtype, device=input_ids.device) # fp16 compatibility | |
batch_size, num_images = pixel_values.size(0), pixel_values.size(1) | |
pixel_values = pixel_values.contiguous().view(batch_size * num_images, *pixel_values.shape[2:]) | |
# Remove padding images - padding images are full 0. | |
real_images_inds = pixel_values.sum(dim=(-1, -2, -3)) != 0.0 | |
pixel_values = pixel_values[real_images_inds] | |
# Get sequence from the vision encoder | |
image_hidden_states = self.vision_model(pixel_values=pixel_values).last_hidden_state | |
# Modality projection | |
image_hidden_states = self.modality_projection(image_hidden_states) | |
if self.config.use_resampler: | |
image_hidden_states = self.perceiver_resampler(image_hidden_states) | |
elif image_hidden_states is not None: | |
image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=input_ids.device) | |
if past_key_values is None: | |
# When we generate, we don't want to replace the potential image_token_id that we generated by images | |
# that simply don't exist | |
new_inp = self.inputs_merger( | |
input_ids=input_ids, | |
inputs_embeds=inputs_embeds, | |
image_hidden_states=image_hidden_states, | |
) | |
inputs_embeds = new_inp["inputs_embeds"] | |
# Can do add some token types embeddings here (image token vs text token) | |
# something like inputs_embeds += self.token_types(token_types) | |
# embed positions | |
if ( | |
attention_mask is not None | |
and hasattr(self.config, "_flash_attn_2_enabled") | |
and self.config._flash_attn_2_enabled | |
and past_key_values is not None | |
): | |
is_padding_right = attention_mask[:, -1].sum().item() != batch_size | |
if is_padding_right: | |
raise ValueError( | |
"You are attempting to perform batched generation with padding_side='right'" | |
" this may lead to unexpected behaviour for Flash Attention version of Mistral. Make sure to " | |
" call `tokenizer.padding_side = 'left'` before tokenizing the input. " | |
) | |
if getattr(self.config, "_flash_attn_2_enabled", False): | |
# 2d mask is passed through the layers | |
attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None | |
else: | |
# 4d mask is passed through the layers | |
attention_mask = _prepare_4d_causal_attention_mask( | |
attention_mask, | |
(batch_size, seq_length), | |
inputs_embeds, | |
past_key_values_length, | |
sliding_window=self.config.sliding_window, | |
) | |
attention_mask[attention_mask == -float("inf")] = torch.finfo(self.dtype).min | |
hidden_states = inputs_embeds | |
if self.gradient_checkpointing and self.training: | |
if use_cache: | |
logger.warning_once( | |
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." | |
) | |
use_cache = False | |
# decoder layers | |
all_hidden_states = () if output_hidden_states else None | |
all_self_attns = () if output_attentions else None | |
next_decoder_cache = () if use_cache else None | |
for idx, decoder_layer in enumerate(self.layers): | |
if output_hidden_states: | |
all_hidden_states += (hidden_states,) | |
past_key_value = past_key_values[idx] if past_key_values is not None else None | |
if self.gradient_checkpointing and self.training: | |
layer_outputs = self._gradient_checkpointing_func( | |
decoder_layer.__call__, | |
hidden_states, | |
attention_mask, | |
position_ids, | |
past_key_value, | |
output_attentions, | |
use_cache, | |
) | |
else: | |
layer_outputs = decoder_layer( | |
hidden_states, | |
attention_mask=attention_mask, | |
position_ids=position_ids, | |
past_key_value=past_key_value, | |
output_attentions=output_attentions, | |
use_cache=use_cache, | |
) | |
hidden_states = layer_outputs[0] | |
if use_cache: | |
next_decoder_cache += (layer_outputs[2 if output_attentions else 1],) | |
if output_attentions: | |
all_self_attns += (layer_outputs[1],) | |
hidden_states = self.norm(hidden_states) | |
# add hidden states from the last decoder layer | |
if output_hidden_states: | |
all_hidden_states += (hidden_states,) | |
next_cache = next_decoder_cache if use_cache else None | |
if not return_dict: | |
return tuple( | |
v | |
for v in [hidden_states, next_cache, all_hidden_states, all_self_attns, image_hidden_states] | |
if v is not None | |
) | |
return VMistralBaseModelOutputWithPast( | |
last_hidden_state=hidden_states, | |
past_key_values=next_cache, | |
hidden_states=all_hidden_states, | |
attentions=all_self_attns, | |
image_hidden_states=image_hidden_states, | |
) | |
class VMistralForVisionText2Text(VMistralPreTrainedModel): | |
_tied_weights_keys = ["lm_head.weight"] | |
def __init__(self, config, vision_model=None): | |
super().__init__(config) | |
self.model = VMistralModel(config, vision_model=vision_model) | |
self.image_token_id = self.config.image_token_id | |
self.lm_head = DecoupledLinear( | |
in_features=config.hidden_size, | |
out_features=config.vocab_size, | |
out_additional_features=config.additional_vocab_size, | |
bias=False, | |
partially_freeze=config.freeze_lm_head, | |
) | |
# Initialize weights and apply final processing | |
self.post_init() | |
def get_input_embeddings(self): | |
return self.model.embed_tokens | |
def set_input_embeddings(self, value): | |
self.model.embed_tokens = value | |
def get_output_embeddings(self): | |
return self.lm_head | |
def set_output_embeddings(self, new_embeddings): | |
self.lm_head = new_embeddings | |
def set_decoder(self, decoder): | |
self.model = decoder | |
def get_decoder(self): | |
return self.model | |
def tie_weights(self): | |
""" | |
Overwrite `transformers.modeling_utils.PreTrainedModel.tie_weights` to handle the case of DecoupledLinear and DecoupledEmbedding. | |
""" | |
output_embeddings = self.get_output_embeddings() | |
input_embeddings = self.get_input_embeddings() | |
if getattr(self.config, "tie_word_embeddings", True): | |
output_embeddings.weight = input_embeddings.weight | |
if input_embeddings.num_additional_embeddings > 0: | |
assert output_embeddings.out_additional_features == input_embeddings.num_additional_embeddings | |
output_embeddings.additional_fc.weight = input_embeddings.additional_embedding.weight | |
if hasattr(output_embeddings, "out_features") and hasattr(input_embeddings, "num_embeddings"): | |
output_embeddings.out_features = input_embeddings.num_embeddings | |
if hasattr(output_embeddings, "out_additional_features") and hasattr( | |
input_embeddings, "num_additional_embeddings" | |
): | |
output_embeddings.out_additional_features = input_embeddings.num_additional_embeddings | |
def forward( | |
self, | |
input_ids: torch.LongTensor = None, | |
attention_mask: Optional[torch.Tensor] = None, | |
position_ids: Optional[torch.LongTensor] = None, | |
past_key_values: Optional[List[torch.FloatTensor]] = None, | |
inputs_embeds: Optional[torch.FloatTensor] = None, | |
pixel_values: Optional[torch.FloatTensor] = None, | |
image_hidden_states: Optional[torch.FloatTensor] = None, | |
labels: Optional[torch.LongTensor] = None, | |
use_cache: Optional[bool] = None, | |
output_attentions: Optional[bool] = None, | |
output_hidden_states: Optional[bool] = None, | |
return_dict: Optional[bool] = None, | |
) -> Union[Tuple, VMistralCausalLMOutputWithPast]: | |
r""" | |
Args: | |
labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): | |
Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., | |
config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored | |
(masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. | |
Returns: | |
""" | |
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions | |
output_hidden_states = ( | |
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states | |
) | |
return_dict = return_dict if return_dict is not None else self.config.use_return_dict | |
# decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) | |
outputs = self.model( | |
input_ids=input_ids, | |
attention_mask=attention_mask, | |
position_ids=position_ids, | |
past_key_values=past_key_values, | |
inputs_embeds=inputs_embeds, | |
pixel_values=pixel_values, | |
image_hidden_states=image_hidden_states, | |
use_cache=use_cache, | |
output_attentions=output_attentions, | |
output_hidden_states=output_hidden_states, | |
return_dict=return_dict, | |
) | |
hidden_states = outputs[0] | |
logits = self.lm_head(hidden_states) | |
logits = logits.float() | |
loss = None | |
if labels is not None: | |
labels = labels.to(logits.device) | |
# Shift so that tokens < n predict n | |
if attention_mask is not None: | |
shift_attention_mask = attention_mask[..., 1:].to(logits.device) | |
shift_logits = logits[..., :-1, :][shift_attention_mask != 0].contiguous() | |
shift_labels = labels[..., 1:][shift_attention_mask != 0].contiguous() | |
else: | |
shift_logits = logits[..., :-1, :].contiguous() | |
shift_labels = labels[..., 1:].contiguous() | |
# Flatten the tokens | |
loss_fct = CrossEntropyLoss(ignore_index=self.image_token_id) | |
loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) | |
if not return_dict: | |
output = (logits,) + outputs[1:] | |
return (loss,) + output if loss is not None else output | |
return VMistralCausalLMOutputWithPast( | |
loss=loss, | |
logits=logits, | |
past_key_values=outputs.past_key_values, | |
hidden_states=outputs.hidden_states, | |
attentions=outputs.attentions, | |
image_hidden_states=outputs.image_hidden_states, | |
) | |
def prepare_inputs_for_generation(self, input_ids, past=None, **kwargs): | |
image_hidden_states = kwargs.pop("image_hidden_states", None) | |
if image_hidden_states is not None: | |
kwargs["pixel_values"] = None | |
inputs = prepare_inputs_for_generation(input_ids, past=past, **kwargs) | |
unwanted_kwargs = ["token_type_ids"] | |
for kwarg in unwanted_kwargs: | |
inputs.pop(kwarg, None) | |
return inputs | |
def _expand_inputs_for_generation( | |
*args, | |
**model_kwargs, | |
): | |
return expand_inputs_for_generation(*args, **model_kwargs) | |
def _update_model_kwargs_for_generation(outputs, model_kwargs, is_encoder_decoder): | |
return update_model_kwargs_for_generation(outputs, model_kwargs) | |
def _reorder_cache(past, beam_idx): | |
reordered_past = () | |
for layer_past in past: | |
reordered_past += (tuple(past_state.index_select(0, beam_idx) for past_state in layer_past),) | |
return reordered_past | |