diff --git a/third_party/transformers/src/transformers/models/altclip/__init__.py b/third_party/transformers/src/transformers/models/altclip/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a30de8a2527567b521be3e9b60e99d025571cb18
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/altclip/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_altclip import *
+ from .modeling_altclip import *
+ from .processing_altclip import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/altclip/configuration_altclip.py b/third_party/transformers/src/transformers/models/altclip/configuration_altclip.py
new file mode 100644
index 0000000000000000000000000000000000000000..c1c69f7cde8904311eba1f35dd68f52aab2b9e8c
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/altclip/configuration_altclip.py
@@ -0,0 +1,224 @@
+# Copyright 2022 WenXiang ZhongzhiCheng LedellWu LiuGuang BoWenZhang and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""AltCLIP model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="BAAI/AltCLIP")
+@strict
+class AltCLIPTextConfig(PreTrainedConfig):
+ r"""
+ project_dim (`int`, *optional*, defaults to 768):
+ The dimensions of the teacher model before the mapping layer.
+
+ Examples:
+
+ ```python
+ >>> from transformers import AltCLIPTextModel, AltCLIPTextConfig
+
+ >>> # Initializing a AltCLIPTextConfig with BAAI/AltCLIP style configuration
+ >>> configuration = AltCLIPTextConfig()
+
+ >>> # Initializing a AltCLIPTextModel (with random weights) from the BAAI/AltCLIP style configuration
+ >>> model = AltCLIPTextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "altclip_text_model"
+
+ vocab_size: int = 250002
+ hidden_size: int = 1024
+ num_hidden_layers: int = 24
+ num_attention_heads: int = 16
+ intermediate_size: int = 4096
+ hidden_act: str = "gelu"
+ hidden_dropout_prob: float | int = 0.1
+ attention_probs_dropout_prob: float | int = 0.1
+ max_position_embeddings: int = 514
+ type_vocab_size: int = 1
+ initializer_range: float = 0.02
+ initializer_factor: float = 0.02
+ layer_norm_eps: float = 1e-05
+ pad_token_id: int | None = 1
+ bos_token_id: int | None = 0
+ eos_token_id: int | list[int] | None = 2
+ project_dim: int = 768
+
+
+@auto_docstring(checkpoint="BAAI/AltCLIP")
+@strict
+class AltCLIPVisionConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AltCLIPVisionConfig, AltCLIPVisionModel
+
+ >>> # Initializing a AltCLIPVisionConfig with BAAI/AltCLIP style configuration
+ >>> configuration = AltCLIPVisionConfig()
+
+ >>> # Initializing a AltCLIPVisionModel (with random weights) from the BAAI/AltCLIP style configuration
+ >>> model = AltCLIPVisionModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "altclip_vision_model"
+ base_config_key = "vision_config"
+
+ hidden_size: int = 768
+ intermediate_size: int = 3072
+ projection_dim: int = 512
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 12
+ num_channels: int = 3
+ image_size: int | list[int] | tuple[int, int] = 224
+ patch_size: int | list[int] | tuple[int, int] = 32
+ hidden_act: str = "quick_gelu"
+ layer_norm_eps: float = 1e-5
+ attention_dropout: float | int = 0.0
+ initializer_range: float = 0.02
+ initializer_factor: float = 1.0
+
+
+@auto_docstring(checkpoint="BAAI/AltCLIP")
+@strict
+class AltCLIPConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AltCLIPConfig, AltCLIPModel
+
+ >>> # Initializing a AltCLIPConfig with BAAI/AltCLIP style configuration
+ >>> configuration = AltCLIPConfig()
+
+ >>> # Initializing a AltCLIPModel (with random weights) from the BAAI/AltCLIP style configuration
+ >>> model = AltCLIPModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+
+ >>> # We can also initialize a AltCLIPConfig from a AltCLIPTextConfig and a AltCLIPVisionConfig
+
+ >>> # Initializing a AltCLIPText and AltCLIPVision configuration
+ >>> config_text = AltCLIPTextConfig()
+ >>> config_vision = AltCLIPVisionConfig()
+
+ >>> config = AltCLIPConfig(text_config=config_text, vision_config=config_vision)
+ ```"""
+
+ model_type = "altclip"
+ sub_configs = {"text_config": AltCLIPTextConfig, "vision_config": AltCLIPVisionConfig}
+ text_config: dict | PreTrainedConfig | None = None
+ vision_config: dict | PreTrainedConfig | None = None
+ projection_dim: int = 768
+ logit_scale_init_value: float = 2.6592
+ initializer_factor: float = 1.0
+
+ def __post_init__(self, **kwargs):
+ if self.text_config is None:
+ text_config = {}
+ logger.info("`text_config` is `None`. Initializing the `AltCLIPTextConfig` with default values.")
+ elif isinstance(self.text_config, AltCLIPTextConfig):
+ text_config = self.text_config.to_dict()
+ else:
+ text_config = self.text_config
+
+ if self.vision_config is None:
+ vision_config = {}
+ logger.info("`vision_config` is `None`. initializing the `AltCLIPVisionConfig` with default values.")
+ elif isinstance(self.vision_config, AltCLIPVisionConfig):
+ vision_config = self.vision_config.to_dict()
+ else:
+ vision_config = self.vision_config
+
+ # For backward compatibility check keyword args
+ # Instead of simply assigning `[text|vision]_config_dict` to `[text|vision]_config`, we use the values in
+ # `[text|vision]_config_dict` to update the values in `[text|vision]_config`. The values should be same in most
+ # cases, but we don't want to break anything regarding `_config_dict` that existed before commit `8827e1b2`.
+ text_config_dict = kwargs.pop("text_config_dict", None)
+ vision_config_dict = kwargs.pop("vision_config_dict", None)
+
+ if text_config_dict is not None:
+ # This is the complete result when using `text_config_dict`.
+ _text_config_dict = AltCLIPTextConfig(**text_config_dict).to_dict()
+
+ # Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
+ for key, value in _text_config_dict.items():
+ if key in text_config and value != text_config[key] and key != "transformers_version":
+ # If specified in `text_config_dict`
+ if key in text_config_dict:
+ message = (
+ f"`{key}` is found in both `text_config_dict` and `text_config` but with different values. "
+ f'The value `text_config_dict["{key}"]` will be used instead.'
+ )
+ # If inferred from default argument values (just to be super careful)
+ else:
+ message = (
+ f"`text_config_dict` is provided which will be used to initialize `AltCLIPTextConfig`. The "
+ f'value `text_config["{key}"]` will be overridden.'
+ )
+ logger.info(message)
+
+ # Update all values in `text_config` with the ones in `_text_config_dict`.
+ text_config.update(_text_config_dict)
+
+ if vision_config_dict is not None:
+ # This is the complete result when using `vision_config_dict`.
+ _vision_config_dict = AltCLIPVisionConfig(**vision_config_dict).to_dict()
+ # convert keys to string instead of integer
+ if "id2label" in _vision_config_dict:
+ _vision_config_dict["id2label"] = {
+ str(key): value for key, value in _vision_config_dict["id2label"].items()
+ }
+
+ # Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
+ for key, value in _vision_config_dict.items():
+ if key in vision_config and value != vision_config[key] and key != "transformers_version":
+ # If specified in `vision_config_dict`
+ if key in vision_config_dict:
+ message = (
+ f"`{key}` is found in both `vision_config_dict` and `vision_config` but with different "
+ f'values. The value `vision_config_dict["{key}"]` will be used instead.'
+ )
+ # If inferred from default argument values (just to be super careful)
+ else:
+ message = (
+ f"`vision_config_dict` is provided which will be used to initialize `AltCLIPVisionConfig`. "
+ f'The value `vision_config["{key}"]` will be overridden.'
+ )
+ logger.info(message)
+
+ # Update all values in `vision_config` with the ones in `_vision_config_dict`.
+ vision_config.update(_vision_config_dict)
+
+ # Finally we can convert back our unified text/vision configs to `PretrainedConfig`
+ self.text_config = AltCLIPTextConfig(**text_config)
+ self.vision_config = AltCLIPVisionConfig(**vision_config)
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["AltCLIPTextConfig", "AltCLIPVisionConfig", "AltCLIPConfig"]
diff --git a/third_party/transformers/src/transformers/models/altclip/modeling_altclip.py b/third_party/transformers/src/transformers/models/altclip/modeling_altclip.py
new file mode 100644
index 0000000000000000000000000000000000000000..5454d012a5a6b817c10087f327ef66f1805aaf68
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/altclip/modeling_altclip.py
@@ -0,0 +1,1185 @@
+# Copyright 2022 The BAAI Teams Authors and The HuggingFace Inc. team. All rights reserved.
+#
+# 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 AltCLIP model."""
+
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPooling,
+ BaseModelOutputWithPoolingAndProjection,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...pytorch_utils import apply_chunking_to_forward
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_altclip import AltCLIPConfig, AltCLIPTextConfig, AltCLIPVisionConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# contrastive loss function, adapted from
+# https://sachinruk.github.io/blog/pytorch/pytorch%20lightning/loss%20function/gpu/2021/03/07/CLIP.html
+def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:
+ return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device))
+
+
+def clip_loss(similarity: torch.Tensor) -> torch.Tensor:
+ caption_loss = contrastive_loss(similarity)
+ image_loss = contrastive_loss(similarity.t())
+ return (caption_loss + image_loss) / 2.0
+
+
+@dataclass
+@auto_docstring
+# Copied from transformers.models.clip.modeling_clip.CLIPOutput with CLIP->AltCLIP
+class AltCLIPOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`):
+ Contrastive loss for image-text similarity.
+ logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`):
+ The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text
+ similarity scores.
+ logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`):
+ The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image
+ similarity scores.
+ text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):
+ The text embeddings obtained by applying the projection layer to the pooled output of [`AltCLIPTextModel`].
+ image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):
+ The image embeddings obtained by applying the projection layer to the pooled output of [`AltCLIPVisionModel`].
+ text_model_output (`BaseModelOutputWithPooling`):
+ The output of the [`AltCLIPTextModel`].
+ vision_model_output (`BaseModelOutputWithPooling`):
+ The output of the [`AltCLIPVisionModel`].
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits_per_image: torch.FloatTensor | None = None
+ logits_per_text: torch.FloatTensor | None = None
+ text_embeds: torch.FloatTensor | None = None
+ image_embeds: torch.FloatTensor | None = None
+ text_model_output: BaseModelOutputWithPooling = None
+ vision_model_output: BaseModelOutputWithPooling = None
+
+ def to_tuple(self) -> tuple[Any]:
+ return tuple(v.to_tuple() if isinstance(v, ModelOutput) else v for v in self.values())
+
+
+# Copied from transformers.models.roberta.modeling_roberta.RobertaEmbeddings with Roberta->AltRoberta
+class AltRobertaEmbeddings(nn.Module):
+ """Construct the embeddings from word, position and token_type embeddings."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+ self.register_buffer(
+ "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+ )
+ self.register_buffer(
+ "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
+ )
+
+ self.padding_idx = config.pad_token_id
+ self.position_embeddings = nn.Embedding(
+ config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
+ )
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values_length: int = 0,
+ ) -> torch.Tensor:
+ if position_ids is None:
+ if input_ids is not None:
+ # Create the position ids from the input token ids. Any padded tokens remain padded.
+ position_ids = self.create_position_ids_from_input_ids(
+ input_ids, self.padding_idx, past_key_values_length
+ )
+ else:
+ position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds, self.padding_idx)
+
+ if input_ids is not None:
+ input_shape = input_ids.size()
+ else:
+ input_shape = inputs_embeds.size()[:-1]
+
+ batch_size, seq_length = input_shape
+
+ # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
+ # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
+ # issue #5664
+ if token_type_ids is None:
+ if hasattr(self, "token_type_ids"):
+ # NOTE: We assume either pos ids to have bsz == 1 (broadcastable) or bsz == effective bsz (input_shape[0])
+ buffered_token_type_ids = self.token_type_ids.expand(position_ids.shape[0], -1)
+ buffered_token_type_ids = torch.gather(buffered_token_type_ids, dim=1, index=position_ids)
+ token_type_ids = buffered_token_type_ids.expand(batch_size, seq_length)
+ else:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.word_embeddings(input_ids)
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
+ embeddings = inputs_embeds + token_type_embeddings
+
+ position_embeddings = self.position_embeddings(position_ids)
+ embeddings = embeddings + position_embeddings
+
+ embeddings = self.LayerNorm(embeddings)
+ embeddings = self.dropout(embeddings)
+ return embeddings
+
+ @staticmethod
+ def create_position_ids_from_inputs_embeds(inputs_embeds, padding_idx):
+ """
+ We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.
+
+ Args:
+ inputs_embeds: torch.Tensor
+
+ Returns: torch.Tensor
+ """
+ input_shape = inputs_embeds.size()[:-1]
+ sequence_length = input_shape[1]
+
+ position_ids = torch.arange(
+ padding_idx + 1, sequence_length + padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
+ )
+ return position_ids.unsqueeze(0).expand(input_shape)
+
+ @staticmethod
+ def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
+ """
+ Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
+ are ignored. This is modified from fairseq's `utils.make_positions`.
+
+ Args:
+ x: torch.Tensor x:
+
+ Returns: torch.Tensor
+ """
+ # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
+ mask = input_ids.ne(padding_idx).int()
+ incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
+ return incremental_indices.long() + padding_idx
+
+
+class AltRobertaSelfAttention(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+ raise ValueError(
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+ f"heads ({config.num_attention_heads})"
+ )
+
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.attention_head_size)
+
+ query_layer = self.query(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_layer = self.key(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_layer = self.value(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+ attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+ if attention_mask is not None:
+ # Apply the attention mask is (precomputed for all layers in AltRobertaModel forward() function)
+ attention_scores = attention_scores + attention_mask
+
+ # Normalize the attention scores to probabilities.
+ attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+
+ # This is actually dropping out entire tokens to attend to, which might
+ # seem a bit unusual, but is taken from the original Transformer paper.
+ attention_probs = self.dropout(attention_probs)
+
+ context_layer = torch.matmul(attention_probs, value_layer)
+
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+ context_layer = context_layer.view(new_context_layer_shape)
+
+ return context_layer, attention_probs
+
+
+# Copied from transformers.models.roberta.modeling_roberta.RobertaSelfOutput
+class AltRobertaSelfOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+ALT_ROBERTA_SELF_ATTENTION_CLASSES = {
+ "eager": AltRobertaSelfAttention,
+}
+
+
+class AltRobertaAttention(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.self = ALT_ROBERTA_SELF_ATTENTION_CLASSES[config._attn_implementation](config)
+ self.output = AltRobertaSelfOutput(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ attention_output, _ = self.self(
+ hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ attention_output = self.output(attention_output, hidden_states)
+ return attention_output
+
+
+# Copied from transformers.models.roberta.modeling_roberta.RobertaIntermediate with Roberta->AltRoberta
+class AltRobertaIntermediate(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.roberta.modeling_roberta.RobertaOutput
+class AltRobertaOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+# Copied from transformers.models.align.modeling_align.AlignTextLayer with AlignText->AltRoberta
+class AltRobertaLayer(GradientCheckpointingLayer):
+ def __init__(self, config):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+ self.attention = AltRobertaAttention(config)
+ self.intermediate = AltRobertaIntermediate(config)
+ self.output = AltRobertaOutput(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ hidden_states = self.attention(
+ hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+
+ hidden_states = apply_chunking_to_forward(
+ self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, hidden_states
+ )
+
+ return hidden_states
+
+ def feed_forward_chunk(self, attention_output):
+ intermediate_output = self.intermediate(attention_output)
+ layer_output = self.output(intermediate_output, attention_output)
+ return layer_output
+
+
+# Copied from transformers.models.align.modeling_align.AlignTextEncoder with AlignText->AltRoberta
+class AltRobertaEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList([AltRobertaLayer(config) for i in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutput:
+ for layer_module in self.layer:
+ hidden_states = layer_module(
+ hidden_states,
+ attention_mask,
+ **kwargs,
+ )
+
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+# Copied from transformers.models.roberta.modeling_roberta.RobertaPooler
+class AltRobertaPooler(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.activation = nn.Tanh()
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ # We "pool" the model by simply taking the hidden state corresponding
+ # to the first token.
+ first_token_tensor = hidden_states[:, 0]
+ pooled_output = self.dense(first_token_tensor)
+ pooled_output = self.activation(pooled_output)
+ return pooled_output
+
+
+# Copied from transformers.models.siglip.modeling_siglip.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class AltCLIPAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = self.embed_dim // self.num_heads
+ if self.head_dim * self.num_heads != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+ f" {self.num_heads})."
+ )
+ self.scale = self.head_dim**-0.5
+ self.dropout = config.attention_dropout
+ self.is_causal = False
+
+ self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Input shape: Batch x Time x Channel"""
+
+ batch_size, seq_length, embed_dim = hidden_states.shape
+
+ queries = self.q_proj(hidden_states)
+ keys = self.k_proj(hidden_states)
+ values = self.v_proj(hidden_states)
+
+ queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+ keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+ values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ queries,
+ keys,
+ values,
+ attention_mask,
+ scaling=self.scale,
+ dropout=0.0 if not self.training else self.dropout,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous()
+ attn_output = self.out_proj(attn_output)
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->AltCLIP
+class AltCLIPMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.activation_fn = ACT2FN[config.hidden_act]
+ self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+ self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ return hidden_states
+
+
+class AltCLIPEncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: AltCLIPConfig):
+ super().__init__()
+ self.embed_dim = config.hidden_size
+ self.self_attn = AltCLIPAttention(config)
+ self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+ self.mlp = AltCLIPMLP(config)
+ self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.FloatTensor, torch.Tensor | None]:
+ residual = hidden_states
+
+ hidden_states = self.layer_norm1(hidden_states)
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.layer_norm2(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+class AltCLIPEncoder(nn.Module):
+ """
+ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+ [`AltCLIPEncoderLayer`].
+
+ Args:
+ config: AltCLIPConfig
+ """
+
+ def __init__(self, config: AltCLIPConfig):
+ super().__init__()
+ self.config = config
+ self.layers = nn.ModuleList([AltCLIPEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ inputs_embeds,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutput:
+ r"""
+ Args:
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ 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.
+ 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)
+
+ """
+ hidden_states = inputs_embeds
+ for encoder_layer in self.layers:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask,
+ **kwargs,
+ )
+
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPVisionEmbeddings with CLIP->AltCLIP
+class AltCLIPVisionEmbeddings(nn.Module):
+ def __init__(self, config: AltCLIPVisionConfig):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.image_size = config.image_size
+ self.patch_size = config.patch_size
+
+ self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))
+
+ self.patch_embedding = nn.Conv2d(
+ in_channels=config.num_channels,
+ out_channels=self.embed_dim,
+ kernel_size=self.patch_size,
+ stride=self.patch_size,
+ bias=False,
+ )
+
+ self.num_patches = (self.image_size // self.patch_size) ** 2
+ self.num_positions = self.num_patches + 1
+ self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+ self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)
+
+ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+ """
+ This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution
+ images. This method is also adapted to support torch.jit tracing.
+
+ Adapted from:
+ - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
+ - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
+ """
+
+ num_patches = embeddings.shape[1] - 1
+ position_embedding = self.position_embedding.weight.unsqueeze(0)
+ num_positions = position_embedding.shape[1] - 1
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embedding(self.position_ids)
+
+ class_pos_embed = position_embedding[:, :1]
+ patch_pos_embed = position_embedding[:, 1:]
+
+ dim = embeddings.shape[-1]
+
+ new_height = height // self.patch_size
+ new_width = width // self.patch_size
+
+ sqrt_num_positions = torch_int(num_positions**0.5)
+ patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
+ patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed,
+ size=(new_height, new_width),
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+ return torch.cat((class_pos_embed, patch_pos_embed), dim=1)
+
+ def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=False) -> torch.Tensor:
+ batch_size, _, height, width = pixel_values.shape
+ if not interpolate_pos_encoding and (height != self.image_size or width != self.image_size):
+ raise ValueError(
+ f"Input image size ({height}*{width}) doesn't match model ({self.image_size}*{self.image_size})."
+ )
+ target_dtype = self.patch_embedding.weight.dtype
+ patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid]
+ patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+
+ class_embeds = self.class_embedding.expand(batch_size, 1, -1)
+ embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
+ if interpolate_pos_encoding:
+ embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+ else:
+ embeddings = embeddings + self.position_embedding(self.position_ids)
+ return embeddings
+
+
+@auto_docstring
+class AltCLIPPreTrainedModel(PreTrainedModel):
+ config: AltCLIPConfig
+ base_model_prefix = "altclip"
+ input_modalities = ("image", "text")
+ supports_gradient_checkpointing = True
+ _no_split_module = []
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ factor = self.config.initializer_factor
+ if isinstance(module, AltCLIPVisionEmbeddings):
+ factor = self.config.initializer_factor
+ init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor)
+ init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor)
+ init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor)
+ init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1)))
+ elif isinstance(module, AltCLIPAttention):
+ factor = self.config.initializer_factor
+ in_proj_std = (module.embed_dim**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
+ out_proj_std = (module.embed_dim**-0.5) * factor
+ init.normal_(module.q_proj.weight, std=in_proj_std)
+ init.normal_(module.k_proj.weight, std=in_proj_std)
+ init.normal_(module.v_proj.weight, std=in_proj_std)
+ init.normal_(module.out_proj.weight, std=out_proj_std)
+ elif isinstance(module, AltCLIPMLP):
+ factor = self.config.initializer_factor
+ in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
+ fc_std = (2 * module.config.hidden_size) ** -0.5 * factor
+ init.normal_(module.fc1.weight, std=fc_std)
+ init.normal_(module.fc2.weight, std=in_proj_std)
+ elif isinstance(module, AltCLIPModel):
+ init.normal_(
+ module.text_projection.weight,
+ std=module.text_embed_dim**-0.5 * self.config.initializer_factor,
+ )
+ init.normal_(
+ module.visual_projection.weight,
+ std=module.vision_embed_dim**-0.5 * self.config.initializer_factor,
+ )
+ elif isinstance(module, nn.LayerNorm):
+ init.zeros_(module.bias)
+ init.ones_(module.weight)
+ elif isinstance(module, nn.Linear):
+ init.normal_(module.weight, mean=0.0, std=self.config.initializer_factor)
+ if module.bias is not None:
+ init.zeros_(module.bias)
+ elif isinstance(module, nn.Embedding):
+ init.normal_(module.weight, mean=0.0, std=self.config.initializer_factor)
+ # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag
+ if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False):
+ init.zeros_(module.weight[module.padding_idx])
+ elif isinstance(module, AltRobertaEmbeddings):
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+ init.zeros_(module.token_type_ids)
+
+
+class AltCLIPVisionTransformer(nn.Module):
+ def __init__(self, config: AltCLIPVisionConfig):
+ super().__init__()
+ self.config = config
+ embed_dim = config.hidden_size
+
+ self.embeddings = AltCLIPVisionEmbeddings(config)
+ self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+ self.encoder = AltCLIPEncoder(config)
+ self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ interpolate_pos_encoding: bool | None = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
+ hidden_states = self.pre_layrnorm(hidden_states)
+
+ encoder_outputs = self.encoder(
+ inputs_embeds=hidden_states,
+ **kwargs,
+ )
+
+ last_hidden_state = encoder_outputs[0]
+ pooled_output = last_hidden_state[:, 0, :]
+ pooled_output = self.post_layernorm(pooled_output)
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooled_output,
+ )
+
+
+class AltCLIPVisionModel(AltCLIPPreTrainedModel):
+ config: AltCLIPVisionConfig
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+ _can_record_outputs = {
+ "hidden_states": AltCLIPEncoderLayer,
+ "attentions": AltCLIPAttention,
+ }
+
+ def __init__(self, config: AltCLIPVisionConfig):
+ super().__init__(config)
+ self.vision_model = AltCLIPVisionTransformer(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.vision_model.embeddings.patch_embedding
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPooling:
+ r"""
+ Examples:
+
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoProcessor, AltCLIPVisionModel
+
+ >>> model = AltCLIPVisionModel.from_pretrained("BAAI/AltCLIP")
+ >>> processor = AutoProcessor.from_pretrained("BAAI/AltCLIP")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> inputs = processor(images=image, return_tensors="pt")
+
+ >>> outputs = model(**inputs)
+ >>> last_hidden_state = outputs.last_hidden_state
+ >>> pooled_output = outputs.pooler_output # pooled CLS states
+ ```"""
+ return self.vision_model(
+ pixel_values=pixel_values,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ **kwargs,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The model behaves as an encoder following the architecture described in *Attention is
+ all you need*_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
+ Kaiser and Illia Polosukhin.
+
+ .. _*Attention is all you need*: https://huggingface.co/papers/1706.03762
+ """
+)
+class AltRobertaModel(AltCLIPPreTrainedModel):
+ config: AltCLIPTextConfig
+ _can_record_outputs = {
+ "hidden_states": AltRobertaLayer,
+ "attentions": AltRobertaSelfAttention,
+ }
+
+ # Copied from transformers.models.clap.modeling_clap.ClapTextModel.__init__ with ClapText->AltRoberta
+ def __init__(self, config, add_pooling_layer=True):
+ r"""
+ add_pooling_layer (bool, *optional*, defaults to `True`):
+ Whether to add a pooling layer
+ """
+ super().__init__(config)
+ self.config = config
+
+ self.embeddings = AltRobertaEmbeddings(config)
+ self.encoder = AltRobertaEncoder(config)
+
+ self.pooler = AltRobertaPooler(config) if add_pooling_layer else None
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embeddings.word_embeddings
+
+ def set_input_embeddings(self, value):
+ self.embeddings.word_embeddings = value
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ # Copied from transformers.models.clap.modeling_clap.ClapTextModel.forward
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPooling:
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+ elif input_ids is not None:
+ self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
+ input_shape = input_ids.size()
+ elif inputs_embeds is not None:
+ input_shape = inputs_embeds.size()[:-1]
+ else:
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+ batch_size, seq_length = input_shape
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+ if attention_mask is None:
+ attention_mask = torch.ones(((batch_size, seq_length)), device=device)
+
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+ # ourselves in which case we just need to make it broadcastable to all heads.
+ extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)
+
+ embedding_output = self.embeddings(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ )
+ encoder_outputs = self.encoder(
+ embedding_output,
+ attention_mask=extended_attention_mask,
+ **kwargs,
+ )
+ sequence_output = encoder_outputs[0]
+ pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ )
+
+
+class AltCLIPTextModel(AltCLIPPreTrainedModel):
+ config: AltCLIPTextConfig
+ input_modalities = ("text",)
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.roberta = AltRobertaModel(config, add_pooling_layer=False)
+ self.transformation = nn.Linear(config.hidden_size, config.project_dim)
+ self.pre_LN = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.roberta.embeddings.word_embeddings
+
+ def set_input_embeddings(self, value: nn.Embedding) -> None:
+ self.roberta.embeddings.word_embeddings = value
+
+ def resize_token_embeddings(self, new_num_tokens: int | None = None) -> nn.Embedding:
+ return super().resize_token_embeddings(new_num_tokens)
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPoolingAndProjection:
+ r"""
+ Examples:
+
+ ```python
+ >>> from transformers import AutoProcessor, AltCLIPTextModel
+
+ >>> model = AltCLIPTextModel.from_pretrained("BAAI/AltCLIP")
+ >>> processor = AutoProcessor.from_pretrained("BAAI/AltCLIP")
+
+ >>> texts = ["it's a cat", "it's a dog"]
+
+ >>> inputs = processor(text=texts, padding=True, return_tensors="pt")
+
+ >>> outputs = model(**inputs)
+ >>> last_hidden_state = outputs.last_hidden_state
+ >>> pooled_output = outputs.pooler_output # pooled CLS states
+ ```"""
+
+ outputs = self.roberta(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ # last module outputs
+ sequence_output = outputs[0]
+
+ # project every module
+ sequence_output = self.pre_LN(sequence_output)
+
+ # pooler
+ projection_state = self.transformation(sequence_output)
+ pooler_output = projection_state[:, 0]
+
+ return BaseModelOutputWithPoolingAndProjection(
+ last_hidden_state=projection_state,
+ pooler_output=pooler_output,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+class AltCLIPModel(AltCLIPPreTrainedModel):
+ config: AltCLIPConfig
+ _can_record_outputs = {
+ "hidden_states": AltCLIPEncoderLayer,
+ "attentions": AltCLIPAttention,
+ }
+
+ def __init__(self, config: AltCLIPConfig):
+ super().__init__(config)
+
+ if not isinstance(config.vision_config, AltCLIPVisionConfig):
+ raise TypeError(
+ "config.vision_config is expected to be of type AltCLIPVisionConfig but is of type"
+ f" {type(config.vision_config)}."
+ )
+ if not isinstance(config.text_config, AltCLIPTextConfig):
+ raise TypeError(
+ "config.text_config is expected to be of type AltCLIPTextConfig but is of type"
+ f" {type(config.text_config)}."
+ )
+
+ text_config = config.text_config
+ vision_config = config.vision_config
+ # The module using it is not a PreTrainedModel subclass so we need this
+ vision_config._attn_implementation = config._attn_implementation
+
+ self.projection_dim = config.projection_dim
+ self.text_embed_dim = text_config.project_dim
+ self.vision_embed_dim = vision_config.hidden_size
+
+ self.text_model = AltCLIPTextModel(text_config)
+ self.vision_model = AltCLIPVisionTransformer(vision_config)
+
+ self.visual_projection = nn.Linear(self.vision_embed_dim, self.projection_dim, bias=False)
+ self.text_projection = nn.Linear(self.text_embed_dim, self.projection_dim, bias=False)
+ self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def get_text_features(
+ self,
+ input_ids: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ Examples:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoProcessor, AltCLIPModel
+
+ >>> model = AltCLIPModel.from_pretrained("BAAI/AltCLIP")
+ >>> processor = AutoProcessor.from_pretrained("BAAI/AltCLIP")
+
+ >>> inputs = processor(text=["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")
+ >>> with torch.inference_mode():
+ ... text_features = model.get_text_features(**inputs)
+ ```"""
+ text_outputs: BaseModelOutputWithPoolingAndProjection = self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ token_type_ids=token_type_ids,
+ **kwargs,
+ )
+ pooled_output = text_outputs.pooler_output
+ text_outputs.pooler_output = self.text_projection(pooled_output)
+
+ return text_outputs
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ interpolate_pos_encoding: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ Examples:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoProcessor, AltCLIPModel
+ >>> from transformers.image_utils import load_image
+
+ >>> model = AltCLIPModel.from_pretrained("BAAI/AltCLIP")
+ >>> processor = AutoProcessor.from_pretrained("BAAI/AltCLIP")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = load_image(url)
+
+ >>> inputs = processor(images=image, return_tensors="pt")
+ >>> with torch.inference_mode():
+ ... image_features = model.get_image_features(**inputs)
+ ```"""
+ vision_outputs = self.vision_model(
+ pixel_values=pixel_values,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ **kwargs,
+ )
+ pooled_output = vision_outputs.pooler_output
+ vision_outputs.pooler_output = self.visual_projection(pooled_output)
+
+ return vision_outputs
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ return_loss: bool | None = None,
+ interpolate_pos_encoding: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | AltCLIPOutput:
+ r"""
+ return_loss (`bool`, *optional*):
+ Whether or not to return the contrastive loss.
+
+ Examples:
+
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoProcessor, AltCLIPModel
+
+ >>> model = AltCLIPModel.from_pretrained("BAAI/AltCLIP")
+ >>> processor = AutoProcessor.from_pretrained("BAAI/AltCLIP")
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+ >>> inputs = processor(
+ ... text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True
+ ... )
+ >>> outputs = model(**inputs)
+ >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score
+ >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities
+ ```"""
+ text_outputs = self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+ vision_outputs = self.vision_model(
+ pixel_values=pixel_values,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ **kwargs,
+ )
+
+ image_embeds = vision_outputs[1]
+ image_embeds = self.visual_projection(image_embeds)
+
+ text_embeds = text_outputs[1]
+ text_embeds = self.text_projection(text_embeds)
+
+ # normalized features
+ image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
+ text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
+
+ # cosine similarity as logits
+ logit_scale = self.logit_scale.exp()
+ logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
+ logits_per_image = logits_per_text.T
+
+ loss = None
+ if return_loss:
+ loss = clip_loss(logits_per_text)
+
+ return AltCLIPOutput(
+ loss=loss,
+ logits_per_image=logits_per_image,
+ logits_per_text=logits_per_text,
+ text_embeds=text_embeds,
+ image_embeds=image_embeds,
+ text_model_output=text_outputs,
+ vision_model_output=vision_outputs,
+ )
+
+
+__all__ = ["AltCLIPPreTrainedModel", "AltCLIPVisionModel", "AltCLIPTextModel", "AltCLIPModel"]
diff --git a/third_party/transformers/src/transformers/models/altclip/processing_altclip.py b/third_party/transformers/src/transformers/models/altclip/processing_altclip.py
new file mode 100644
index 0000000000000000000000000000000000000000..78fd31be5f227bfc9eef107ce74c0c88ff0dbbec
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/altclip/processing_altclip.py
@@ -0,0 +1,28 @@
+# Copyright 2022 WenXiang ZhongzhiCheng LedellWu LiuGuang BoWenZhang The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""
+Image/Text processor class for AltCLIP
+"""
+
+from ...processing_utils import ProcessorMixin
+from ...utils import auto_docstring
+
+
+@auto_docstring
+class AltCLIPProcessor(ProcessorMixin):
+ def __init__(self, image_processor=None, tokenizer=None):
+ super().__init__(image_processor, tokenizer)
+
+
+__all__ = ["AltCLIPProcessor"]
diff --git a/third_party/transformers/src/transformers/models/auto/__init__.py b/third_party/transformers/src/transformers/models/auto/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6b86884b3b7b0819cc58157a6593aa5d53c883b9
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/auto/__init__.py
@@ -0,0 +1,33 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .auto_factory import *
+ from .configuration_auto import *
+ from .feature_extraction_auto import *
+ from .image_processing_auto import *
+ from .modeling_auto import *
+ from .processing_auto import *
+ from .tokenization_auto import *
+ from .video_processing_auto import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/auto/auto_factory.py b/third_party/transformers/src/transformers/models/auto/auto_factory.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f93efb1f57e428574e1ee98182bc205488e47e4
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/auto/auto_factory.py
@@ -0,0 +1,662 @@
+# Copyright 2021 The HuggingFace Inc. team.
+#
+# 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.
+"""Factory function to build auto-model classes."""
+
+import copy
+import importlib
+import json
+import os
+from collections import OrderedDict
+from collections.abc import Iterator
+from typing import Any, TypeVar
+
+from huggingface_hub import repo_exists
+
+from ...configuration_utils import PreTrainedConfig
+from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
+from ...utils import (
+ CONFIG_NAME,
+ cached_file,
+ copy_func,
+ extract_commit_hash,
+ find_adapter_config_file,
+ is_peft_available,
+ is_torch_available,
+ logging,
+ requires_backends,
+)
+from .configuration_auto import AutoConfig, model_type_to_module_name, replace_list_option_in_docstrings
+
+
+if is_torch_available():
+ from ...generation import GenerationMixin
+
+
+logger = logging.get_logger(__name__)
+
+_T = TypeVar("_T")
+# Tokenizers will depend on packages installed, too much variance and there are no common base or Protocol
+_LazyAutoMappingValue = tuple[type[Any] | None, type[Any] | None]
+
+CLASS_DOCSTRING = """
+ This is a generic model class that will be instantiated as one of the model classes of the library when created
+ with the [`~BaseAutoModelClass.from_pretrained`] class method or the [`~BaseAutoModelClass.from_config`] class
+ method.
+
+ This class cannot be instantiated directly using `__init__()` (throws an error).
+"""
+
+FROM_CONFIG_DOCSTRING = """
+ Instantiates one of the model classes of the library from a configuration.
+
+ Note:
+ Loading a model from its configuration file does **not** load the model weights. It only affects the
+ model's configuration. Use [`~BaseAutoModelClass.from_pretrained`] to load the model weights.
+
+ Args:
+ config ([`PreTrainedConfig`]):
+ The model class to instantiate is selected based on the configuration class:
+
+ List options
+ attn_implementation (`str`, *optional*):
+ The attention implementation to use in the model (if relevant). Can be any of `"eager"` (manual implementation of the attention), `"sdpa"` (using [`F.scaled_dot_product_attention`](https://pytorch.org/docs/master/generated/torch.nn.functional.scaled_dot_product_attention.html)), `"flash_attention_2"` (using [Dao-AILab/flash-attention](https://github.com/Dao-AILab/flash-attention)), or `"flash_attention_3"` (using [Dao-AILab/flash-attention/hopper](https://github.com/Dao-AILab/flash-attention/tree/main/hopper)). By default, if available, SDPA will be used for torch>=2.1.1. The default is otherwise the manual `"eager"` implementation.
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoConfig, BaseAutoModelClass
+
+ >>> # Download configuration from huggingface.co and cache.
+ >>> config = AutoConfig.from_pretrained("checkpoint_placeholder")
+ >>> model = BaseAutoModelClass.from_config(config)
+ ```
+"""
+
+FROM_PRETRAINED_TORCH_DOCSTRING = """
+ Instantiate one of the model classes of the library from a pretrained model.
+
+ The model class to instantiate is selected based on the `model_type` property of the config object (either
+ passed as an argument or loaded from `pretrained_model_name_or_path` if possible), or when it's missing, by
+ falling back to using pattern matching on `pretrained_model_name_or_path`:
+
+ List options
+
+ The model is set in evaluation mode by default using `model.eval()` (so for instance, dropout modules are
+ deactivated). To train the model, you should first set it back in training mode with `model.train()`
+
+ Args:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ Can be either:
+
+ - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
+ - A path to a *directory* containing model weights saved using
+ [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
+ model_args (additional positional arguments, *optional*):
+ Will be passed along to the underlying model `__init__()` method.
+ config ([`PreTrainedConfig`], *optional*):
+ Configuration for the model to use instead of an automatically loaded configuration. Configuration can
+ be automatically loaded when:
+
+ - The model is a model provided by the library (loaded with the *model id* string of a pretrained
+ model).
+ - The model was saved using [`~PreTrainedModel.save_pretrained`] and is reloaded by supplying the
+ save directory.
+ - The model is loaded by supplying a local directory as `pretrained_model_name_or_path` and a
+ configuration JSON file named *config.json* is found in the directory.
+ state_dict (*dict[str, torch.Tensor]*, *optional*):
+ A state dictionary to use instead of a state dictionary loaded from saved weights file.
+
+ This option can be used if you want to create a model from a pretrained configuration but load your own
+ weights. In this case though, you should check if using [`~PreTrainedModel.save_pretrained`] and
+ [`~PreTrainedModel.from_pretrained`] is not a simpler option.
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the
+ standard cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force the (re-)download of the model weights and configuration files, overriding the
+ cached versions if they exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
+ output_loading_info(`bool`, *optional*, defaults to `False`):
+ Whether or not to also return a dictionary containing missing keys, unexpected keys and error messages.
+ local_files_only(`bool`, *optional*, defaults to `False`):
+ Whether or not to only look at local files (e.g., not try downloading the model).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ trust_remote_code (`bool`, *optional*, defaults to `False`):
+ Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
+ should only be set to `True` for repositories you trust and in which you have read the code, as it will
+ execute code present on the Hub on your local machine.
+ code_revision (`str`, *optional*, defaults to `"main"`):
+ The specific revision to use for the code on the Hub, if the code leaves in a different repository than
+ the rest of the model. It can be a branch name, a tag name, or a commit id, since we use a git-based
+ system for storing models and other artifacts on huggingface.co, so `revision` can be any identifier
+ allowed by git.
+ kwargs (additional keyword arguments, *optional*):
+ Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
+ `output_attentions=True`). Behaves differently depending on whether a `config` is provided or
+ automatically loaded:
+
+ - If a configuration is provided with `config`, `**kwargs` will be directly passed to the
+ underlying model's `__init__` method (we assume all relevant updates to the configuration have
+ already been done)
+ - If a configuration is not provided, `kwargs` will be first passed to the configuration class
+ initialization function ([`~PreTrainedConfig.from_pretrained`]). Each key of `kwargs` that
+ corresponds to a configuration attribute will be used to override said attribute with the
+ supplied `kwargs` value. Remaining keys that do not correspond to any configuration attribute
+ will be passed to the underlying model's `__init__` function.
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoConfig, BaseAutoModelClass
+
+ >>> # Download model and configuration from huggingface.co and cache.
+ >>> model = BaseAutoModelClass.from_pretrained("checkpoint_placeholder")
+
+ >>> # Update configuration during loading
+ >>> model = BaseAutoModelClass.from_pretrained("checkpoint_placeholder", output_attentions=True)
+ >>> model.config.output_attentions
+ True
+ ```
+"""
+
+
+def _get_model_class(config, model_mapping):
+ supported_models = model_mapping[type(config)]
+ if not isinstance(supported_models, (list, tuple)):
+ return supported_models
+
+ name_to_model = {model.__name__: model for model in supported_models}
+ architectures = getattr(config, "architectures", [])
+ for arch in architectures:
+ if arch in name_to_model:
+ return name_to_model[arch]
+
+ # If not architecture is set in the config or match the supported models, the first element of the tuple is the
+ # defaults.
+ return supported_models[0]
+
+
+class _BaseAutoModelClass:
+ # Base class for auto models.
+ _model_mapping = None
+
+ def __init__(self, *args, **kwargs) -> None:
+ raise OSError(
+ f"{self.__class__.__name__} is designed to be instantiated "
+ f"using the `{self.__class__.__name__}.from_pretrained(pretrained_model_name_or_path)` or "
+ f"`{self.__class__.__name__}.from_config(config)` methods."
+ )
+
+ @classmethod
+ def from_config(cls, config, **kwargs):
+ trust_remote_code = kwargs.pop("trust_remote_code", None)
+ has_remote_code = hasattr(config, "auto_map") and cls.__name__ in config.auto_map
+ has_local_code = type(config) in cls._model_mapping
+ explicit_local_code = has_local_code and not _get_model_class(
+ config, cls._model_mapping
+ ).__module__.startswith("transformers.")
+ if has_remote_code:
+ class_ref = config.auto_map[cls.__name__]
+ if "--" in class_ref:
+ upstream_repo = class_ref.split("--")[0]
+ else:
+ upstream_repo = None
+ trust_remote_code = resolve_trust_remote_code(
+ trust_remote_code, config._name_or_path, has_local_code, has_remote_code, upstream_repo=upstream_repo
+ )
+
+ if has_remote_code and trust_remote_code and not explicit_local_code:
+ if "--" in class_ref:
+ repo_id, class_ref = class_ref.split("--")
+ else:
+ repo_id = config.name_or_path
+ model_class = get_class_from_dynamic_module(class_ref, repo_id, **kwargs)
+ # This block handles the case where the user is loading a model with `trust_remote_code=True`
+ # but a library model exists with the same name. We don't want to override the autoclass
+ # mappings in this case, or all future loads of that model will be the remote code model.
+ if not has_local_code:
+ cls.register(config.__class__, model_class, exist_ok=True)
+ model_class.register_for_auto_class(auto_class=cls)
+ _ = kwargs.pop("code_revision", None)
+ model_class = add_generation_mixin_to_remote_model(model_class)
+ return model_class._from_config(config, **kwargs)
+ elif has_local_code:
+ model_class = _get_model_class(config, cls._model_mapping)
+ return model_class._from_config(config, **kwargs)
+
+ raise ValueError(
+ f"Unrecognized configuration class {config.__class__} for this kind of AutoModel: {cls.__name__}.\n"
+ f"Model type should be one of {', '.join(c.__name__ for c in cls._model_mapping)}."
+ )
+
+ @classmethod
+ def _prepare_config_for_auto_class(cls, config: PreTrainedConfig) -> PreTrainedConfig:
+ """Additional autoclass-specific config post-loading manipulation. May be overridden in subclasses."""
+ return config
+
+ @classmethod
+ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike[str], *model_args, **kwargs):
+ config = kwargs.pop("config", None)
+ trust_remote_code = kwargs.get("trust_remote_code")
+ kwargs["_from_auto"] = True
+ hub_kwargs_names = [
+ "cache_dir",
+ "force_download",
+ "local_files_only",
+ "proxies",
+ "revision",
+ "subfolder",
+ "token",
+ ]
+ hub_kwargs = {name: kwargs.pop(name) for name in hub_kwargs_names if name in kwargs}
+ code_revision = kwargs.pop("code_revision", None)
+ commit_hash = kwargs.pop("_commit_hash", None)
+ adapter_kwargs = kwargs.pop("adapter_kwargs", None)
+
+ token = hub_kwargs.pop("token", None)
+
+ if token is not None:
+ hub_kwargs["token"] = token
+
+ if commit_hash is None:
+ if not isinstance(config, PreTrainedConfig):
+ # We make a call to the config file first (which may be absent) to get the commit hash as soon as possible
+ resolved_config_file = cached_file(
+ pretrained_model_name_or_path,
+ CONFIG_NAME,
+ _raise_exceptions_for_gated_repo=False,
+ _raise_exceptions_for_missing_entries=False,
+ _raise_exceptions_for_connection_errors=False,
+ **hub_kwargs,
+ )
+ commit_hash = extract_commit_hash(resolved_config_file, commit_hash)
+ else:
+ commit_hash = getattr(config, "_commit_hash", None)
+
+ if is_peft_available():
+ if adapter_kwargs is None:
+ adapter_kwargs = {}
+ adapter_kwargs = adapter_kwargs.copy() # avoid mutating original
+ if token is not None:
+ adapter_kwargs["token"] = token
+
+ maybe_adapter_path = find_adapter_config_file(
+ pretrained_model_name_or_path, _commit_hash=commit_hash, **adapter_kwargs
+ )
+
+ if maybe_adapter_path is not None:
+ with open(maybe_adapter_path, "r", encoding="utf-8") as f:
+ adapter_config = json.load(f)
+
+ adapter_kwargs["_adapter_model_path"] = pretrained_model_name_or_path
+ # Only override the model name/path if the current value doesn't point to a
+ # complete model with an embedded adapter so that local models with embedded
+ # adapters will load from the local base model rather than pull the base
+ # model named in the adapter's config from the hub.
+ if not os.path.exists(pretrained_model_name_or_path) or not os.path.exists(
+ os.path.join(pretrained_model_name_or_path, CONFIG_NAME)
+ ):
+ pretrained_model_name_or_path = adapter_config["base_model_name_or_path"]
+
+ if not isinstance(config, PreTrainedConfig):
+ kwargs_orig = copy.deepcopy(kwargs)
+ # ensure not to pollute the config object with dtype="auto" - since it's
+ # meaningless in the context of the config object - torch.dtype values are acceptable
+ if kwargs.get("torch_dtype") == "auto":
+ _ = kwargs.pop("torch_dtype")
+ if kwargs.get("dtype") == "auto":
+ _ = kwargs.pop("dtype")
+ # to not overwrite the quantization_config if config has a quantization_config
+ if kwargs.get("quantization_config") is not None:
+ _ = kwargs.pop("quantization_config")
+
+ config, kwargs = AutoConfig.from_pretrained(
+ pretrained_model_name_or_path,
+ return_unused_kwargs=True,
+ code_revision=code_revision,
+ _commit_hash=commit_hash,
+ **hub_kwargs,
+ **kwargs,
+ )
+
+ # if torch_dtype=auto was passed here, ensure to pass it on
+ if kwargs_orig.get("torch_dtype", None) == "auto":
+ kwargs["torch_dtype"] = "auto"
+ if kwargs_orig.get("dtype", None) == "auto":
+ kwargs["dtype"] = "auto"
+ if kwargs_orig.get("quantization_config", None) is not None:
+ kwargs["quantization_config"] = kwargs_orig["quantization_config"]
+
+ has_remote_code = hasattr(config, "auto_map") and cls.__name__ in config.auto_map
+ has_local_code = type(config) in cls._model_mapping
+ explicit_local_code = has_local_code and not _get_model_class(
+ config, cls._model_mapping
+ ).__module__.startswith("transformers.")
+ upstream_repo = None
+ if has_remote_code:
+ class_ref = config.auto_map[cls.__name__]
+ if "--" in class_ref:
+ upstream_repo = class_ref.split("--")[0]
+ trust_remote_code = resolve_trust_remote_code(
+ trust_remote_code,
+ pretrained_model_name_or_path,
+ has_local_code,
+ has_remote_code,
+ upstream_repo=upstream_repo,
+ )
+ kwargs["trust_remote_code"] = trust_remote_code
+
+ # Set the adapter kwargs
+ kwargs["adapter_kwargs"] = adapter_kwargs
+
+ if has_remote_code and trust_remote_code and not explicit_local_code:
+ model_class = get_class_from_dynamic_module(
+ class_ref, pretrained_model_name_or_path, code_revision=code_revision, **hub_kwargs, **kwargs
+ )
+ _ = hub_kwargs.pop("code_revision", None)
+ # This block handles the case where the user is loading a model with `trust_remote_code=True`
+ # but a library model exists with the same name. We don't want to override the autoclass
+ # mappings in this case, or all future loads of that model will be the remote code model.
+ if not has_local_code:
+ cls.register(config.__class__, model_class, exist_ok=True)
+ model_class.register_for_auto_class(auto_class=cls)
+ model_class = add_generation_mixin_to_remote_model(model_class)
+ return model_class.from_pretrained(
+ pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs, **kwargs
+ )
+ elif has_local_code:
+ model_class = _get_model_class(config, cls._model_mapping)
+ if model_class.config_class == config.sub_configs.get("text_config", None):
+ config = config.get_text_config()
+ return model_class.from_pretrained(
+ pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs, **kwargs
+ )
+ raise ValueError(
+ f"Unrecognized configuration class {config.__class__} for this kind of AutoModel: {cls.__name__}.\n"
+ f"Model type should be one of {', '.join(c.__name__ for c in cls._model_mapping)}."
+ )
+
+ @classmethod
+ def register(cls, config_class, model_class, exist_ok=False) -> None:
+ """
+ Register a new model for this class.
+
+ Args:
+ config_class ([`PreTrainedConfig`]):
+ The configuration corresponding to the model to register.
+ model_class ([`PreTrainedModel`]):
+ The model to register.
+ """
+ if hasattr(model_class, "config_class") and model_class.config_class.__name__ != config_class.__name__:
+ raise ValueError(
+ "The model class you are passing has a `config_class` attribute that is not consistent with the "
+ f"config class you passed (model has {model_class.config_class} and you passed {config_class}. Fix "
+ "one of those so they match!"
+ )
+ cls._model_mapping.register(config_class, model_class, exist_ok=exist_ok)
+
+
+class _BaseAutoBackboneClass(_BaseAutoModelClass):
+ # Base class for auto backbone models.
+ _model_mapping = None
+
+ @classmethod
+ def _load_timm_backbone_from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+ requires_backends(cls, ["vision", "timm"])
+ from ...models.timm_backbone import TimmBackboneConfig
+
+ config = kwargs.pop("config", TimmBackboneConfig())
+
+ if kwargs.get("out_features") is not None:
+ raise ValueError("Cannot specify `out_features` for timm backbones")
+
+ if kwargs.get("output_loading_info", False):
+ raise ValueError("Cannot specify `output_loading_info=True` when loading from timm")
+
+ num_channels = kwargs.pop("num_channels", config.num_channels)
+ features_only = kwargs.pop("features_only", config.features_only)
+ out_indices = kwargs.pop("out_indices", config.out_indices)
+ config = TimmBackboneConfig(
+ backbone=pretrained_model_name_or_path,
+ num_channels=num_channels,
+ features_only=features_only,
+ out_indices=out_indices,
+ )
+ # Always load a pretrained model when `from_pretrained` is called
+ kwargs.pop("use_pretrained_backbone", None)
+ return super().from_config(config, pretrained=True, **kwargs)
+
+ @classmethod
+ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+ kwargs.pop("use_timm_backbone", None)
+ if not repo_exists(pretrained_model_name_or_path):
+ return cls._load_timm_backbone_from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+
+ return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+
+
+def insert_head_doc(docstring, head_doc: str = ""):
+ if len(head_doc) > 0:
+ return docstring.replace(
+ "one of the model classes of the library ",
+ f"one of the model classes of the library (with a {head_doc} head) ",
+ )
+ return docstring.replace(
+ "one of the model classes of the library ", "one of the base model classes of the library "
+ )
+
+
+def auto_class_update(cls, checkpoint_for_example: str = "google-bert/bert-base-cased", head_doc: str = ""):
+ # Create a new class with the right name from the base class
+ model_mapping = cls._model_mapping
+ name = cls.__name__
+ class_docstring = insert_head_doc(CLASS_DOCSTRING, head_doc=head_doc)
+ cls.__doc__ = class_docstring.replace("BaseAutoModelClass", name)
+
+ # Now we need to copy and re-register `from_config` and `from_pretrained` as class methods otherwise we can't
+ # have a specific docstrings for them.
+ from_config = copy_func(_BaseAutoModelClass.from_config)
+ from_config_docstring = insert_head_doc(FROM_CONFIG_DOCSTRING, head_doc=head_doc)
+ from_config_docstring = from_config_docstring.replace("BaseAutoModelClass", name)
+ from_config_docstring = from_config_docstring.replace("checkpoint_placeholder", checkpoint_for_example)
+ from_config.__doc__ = from_config_docstring
+ from_config = replace_list_option_in_docstrings(model_mapping._model_mapping, use_model_types=False)(from_config)
+ cls.from_config = classmethod(from_config)
+
+ from_pretrained_docstring = FROM_PRETRAINED_TORCH_DOCSTRING
+ from_pretrained = copy_func(_BaseAutoModelClass.from_pretrained)
+ from_pretrained_docstring = insert_head_doc(from_pretrained_docstring, head_doc=head_doc)
+ from_pretrained_docstring = from_pretrained_docstring.replace("BaseAutoModelClass", name)
+ from_pretrained_docstring = from_pretrained_docstring.replace("checkpoint_placeholder", checkpoint_for_example)
+ shortcut = checkpoint_for_example.split("/")[-1].split("-")[0]
+ from_pretrained_docstring = from_pretrained_docstring.replace("shortcut_placeholder", shortcut)
+ from_pretrained.__doc__ = from_pretrained_docstring
+ from_pretrained = replace_list_option_in_docstrings(model_mapping._model_mapping)(from_pretrained)
+ cls.from_pretrained = classmethod(from_pretrained)
+ return cls
+
+
+def get_values(model_mapping):
+ result = []
+ for model in model_mapping.values():
+ if isinstance(model, (list, tuple)):
+ result += list(model)
+ else:
+ result.append(model)
+
+ return result
+
+
+def getattribute_from_module(module, attr):
+ if attr is None:
+ return None
+ if isinstance(attr, tuple):
+ return tuple(getattribute_from_module(module, a) for a in attr)
+ if isinstance(attr, dict):
+ return {k: getattribute_from_module(module, v) for k, v in attr.items()}
+ if hasattr(module, attr):
+ return getattr(module, attr)
+ # Some of the mappings have entries model_type -> object of another model type. In that case we try to grab the
+ # object at the top level.
+ transformers_module = importlib.import_module("transformers")
+
+ if module != transformers_module:
+ try:
+ return getattribute_from_module(transformers_module, attr)
+ except ValueError:
+ raise ValueError(f"Could not find {attr} neither in {module} nor in {transformers_module}!")
+ else:
+ raise ValueError(f"Could not find {attr} in {transformers_module}!")
+
+
+def add_generation_mixin_to_remote_model(model_class):
+ """
+ Adds `GenerationMixin` to the inheritance of `model_class`, if `model_class` is a PyTorch model.
+
+ This function is used for backwards compatibility purposes: in v4.45, we've started a deprecation cycle to make
+ `PreTrainedModel` stop inheriting from `GenerationMixin`. Without this function, older models dynamically loaded
+ from the Hub may not have the `generate` method after we remove the inheritance.
+ """
+ # 1. If it is not a PT model (i.e. doesn't inherit Module), do nothing
+ if "torch.nn.modules.module.Module" not in str(model_class.__mro__):
+ return model_class
+
+ # 2. If it already **directly** inherits from GenerationMixin, do nothing
+ if "GenerationMixin" in str(model_class.__bases__):
+ return model_class
+
+ # 3. Prior to v4.45, we could detect whether a model was `generate`-compatible if it had its own `generate` and/or
+ # `prepare_inputs_for_generation` method.
+ has_custom_generate_in_class = hasattr(model_class, "generate") and "GenerationMixin" not in str(
+ getattr(model_class, "generate")
+ )
+ has_custom_prepare_inputs = hasattr(model_class, "prepare_inputs_for_generation") and "GenerationMixin" not in str(
+ getattr(model_class, "prepare_inputs_for_generation")
+ )
+ if has_custom_generate_in_class or has_custom_prepare_inputs:
+ model_class_with_generation_mixin = type(
+ model_class.__name__, (model_class, GenerationMixin), {**model_class.__dict__}
+ )
+ return model_class_with_generation_mixin
+ return model_class
+
+
+class _LazyAutoMapping(OrderedDict[type[PreTrainedConfig], _LazyAutoMappingValue]):
+ """
+ A mapping config to object (model or tokenizer for instance) that will load keys and values when it is accessed.
+
+ Args:
+ - config_mapping: The map model type to config class
+ - model_mapping: The map model type to model (or tokenizer) class
+ """
+
+ def __init__(self, config_mapping, model_mapping) -> None:
+ self._config_mapping = config_mapping
+ self._reverse_config_mapping = {v: k for k, v in config_mapping.items()}
+ self._model_mapping = model_mapping
+ self._model_mapping._model_mapping = self
+ self._extra_content = {}
+ self._modules = {}
+
+ def __len__(self) -> int:
+ common_keys = set(self._config_mapping.keys()).intersection(self._model_mapping.keys())
+ return len(common_keys) + len(self._extra_content)
+
+ def __getitem__(self, key: type[PreTrainedConfig]) -> _LazyAutoMappingValue:
+ if key in self._extra_content:
+ return self._extra_content[key]
+ model_type = self._reverse_config_mapping[key.__name__]
+ if model_type in self._model_mapping:
+ model_name = self._model_mapping[model_type]
+ return self._load_attr_from_module(model_type, model_name)
+
+ # Maybe there was several model types associated with this config.
+ model_types = [k for k, v in self._config_mapping.items() if v == key.__name__]
+ for mtype in model_types:
+ if mtype in self._model_mapping:
+ model_name = self._model_mapping[mtype]
+ return self._load_attr_from_module(mtype, model_name)
+ raise KeyError(key)
+
+ def _load_attr_from_module(self, model_type, attr):
+ module_name = model_type_to_module_name(model_type)
+ if module_name not in self._modules:
+ self._modules[module_name] = importlib.import_module(f".{module_name}", "transformers.models")
+ return getattribute_from_module(self._modules[module_name], attr)
+
+ def keys(self) -> list[type[PreTrainedConfig]]:
+ mapping_keys = [
+ self._load_attr_from_module(key, name)
+ for key, name in self._config_mapping.items()
+ if key in self._model_mapping
+ ]
+ return mapping_keys + list(self._extra_content.keys())
+
+ def get(self, key: type[PreTrainedConfig], default: _T) -> _LazyAutoMappingValue | _T:
+ try:
+ return self.__getitem__(key)
+ except KeyError:
+ return default
+
+ def __bool__(self) -> bool:
+ return bool(self.keys())
+
+ def values(self) -> list[_LazyAutoMappingValue]:
+ mapping_values = [
+ self._load_attr_from_module(key, name)
+ for key, name in self._model_mapping.items()
+ if key in self._config_mapping
+ ]
+ return mapping_values + list(self._extra_content.values())
+
+ def items(self) -> list[tuple[type[PreTrainedConfig], _LazyAutoMappingValue]]:
+ mapping_items = [
+ (
+ self._load_attr_from_module(key, self._config_mapping[key]),
+ self._load_attr_from_module(key, self._model_mapping[key]),
+ )
+ for key in self._model_mapping
+ if key in self._config_mapping
+ ]
+ return mapping_items + list(self._extra_content.items())
+
+ def __iter__(self) -> Iterator[type[PreTrainedConfig]]:
+ return iter(self.keys())
+
+ def __contains__(self, item: type) -> bool:
+ if item in self._extra_content:
+ return True
+ if not hasattr(item, "__name__") or item.__name__ not in self._reverse_config_mapping:
+ return False
+ model_type = self._reverse_config_mapping[item.__name__]
+ return model_type in self._model_mapping
+
+ def register(self, key: type[PreTrainedConfig], value: _LazyAutoMappingValue, exist_ok=False) -> None:
+ """
+ Register a new model in this mapping.
+ """
+ if hasattr(key, "__name__") and key.__name__ in self._reverse_config_mapping:
+ model_type = self._reverse_config_mapping[key.__name__]
+ if model_type in self._model_mapping and not exist_ok:
+ raise ValueError(f"'{key}' is already used by a Transformers model.")
+
+ self._extra_content[key] = value
+
+
+__all__ = ["get_values"]
diff --git a/third_party/transformers/src/transformers/models/auto/configuration_auto.py b/third_party/transformers/src/transformers/models/auto/configuration_auto.py
new file mode 100644
index 0000000000000000000000000000000000000000..2c0fe88d0e74e7be9ece3f459eb6af2c7ff1b270
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/auto/configuration_auto.py
@@ -0,0 +1,1551 @@
+# Copyright 2018 The HuggingFace Inc. team.
+#
+# 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.
+"""Auto Config class."""
+
+import importlib
+import os
+import re
+from collections import OrderedDict
+from collections.abc import Callable, Iterator, KeysView, ValuesView
+from typing import Any, TypeVar
+
+from ...configuration_utils import PreTrainedConfig
+from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
+from ...utils import CONFIG_NAME, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+_CallableT = TypeVar("_CallableT", bound=Callable[..., Any])
+
+
+CONFIG_MAPPING_NAMES = OrderedDict[str, str](
+ [
+ # Add configs here
+ ("afmoe", "AfmoeConfig"),
+ ("aimv2", "Aimv2Config"),
+ ("aimv2_vision_model", "Aimv2VisionConfig"),
+ ("albert", "AlbertConfig"),
+ ("align", "AlignConfig"),
+ ("altclip", "AltCLIPConfig"),
+ ("apertus", "ApertusConfig"),
+ ("arcee", "ArceeConfig"),
+ ("aria", "AriaConfig"),
+ ("aria_text", "AriaTextConfig"),
+ ("audio-spectrogram-transformer", "ASTConfig"),
+ ("audioflamingo3", "AudioFlamingo3Config"),
+ ("audioflamingo3_encoder", "AudioFlamingo3EncoderConfig"),
+ ("autoformer", "AutoformerConfig"),
+ ("aya_vision", "AyaVisionConfig"),
+ ("bamba", "BambaConfig"),
+ ("bark", "BarkConfig"),
+ ("bart", "BartConfig"),
+ ("beit", "BeitConfig"),
+ ("bert", "BertConfig"),
+ ("bert-generation", "BertGenerationConfig"),
+ ("big_bird", "BigBirdConfig"),
+ ("bigbird_pegasus", "BigBirdPegasusConfig"),
+ ("biogpt", "BioGptConfig"),
+ ("bit", "BitConfig"),
+ ("bitnet", "BitNetConfig"),
+ ("blenderbot", "BlenderbotConfig"),
+ ("blenderbot-small", "BlenderbotSmallConfig"),
+ ("blip", "BlipConfig"),
+ ("blip-2", "Blip2Config"),
+ ("blip_2_qformer", "Blip2QFormerConfig"),
+ ("bloom", "BloomConfig"),
+ ("blt", "BltConfig"),
+ ("bridgetower", "BridgeTowerConfig"),
+ ("bros", "BrosConfig"),
+ ("camembert", "CamembertConfig"),
+ ("canine", "CanineConfig"),
+ ("chameleon", "ChameleonConfig"),
+ ("chinese_clip", "ChineseCLIPConfig"),
+ ("chinese_clip_vision_model", "ChineseCLIPVisionConfig"),
+ ("chmv2", "CHMv2Config"),
+ ("clap", "ClapConfig"),
+ ("clip", "CLIPConfig"),
+ ("clip_text_model", "CLIPTextConfig"),
+ ("clip_vision_model", "CLIPVisionConfig"),
+ ("clipseg", "CLIPSegConfig"),
+ ("clvp", "ClvpConfig"),
+ ("code_llama", "LlamaConfig"),
+ ("codegen", "CodeGenConfig"),
+ ("cohere", "CohereConfig"),
+ ("cohere2", "Cohere2Config"),
+ ("cohere2_vision", "Cohere2VisionConfig"),
+ ("cohere_asr", "CohereAsrConfig"),
+ ("colmodernvbert", "ColModernVBertConfig"),
+ ("colpali", "ColPaliConfig"),
+ ("colqwen2", "ColQwen2Config"),
+ ("conditional_detr", "ConditionalDetrConfig"),
+ ("convbert", "ConvBertConfig"),
+ ("convnext", "ConvNextConfig"),
+ ("convnextv2", "ConvNextV2Config"),
+ ("cpmant", "CpmAntConfig"),
+ ("csm", "CsmConfig"),
+ ("ctrl", "CTRLConfig"),
+ ("cvt", "CvtConfig"),
+ ("cwm", "CwmConfig"),
+ ("d_fine", "DFineConfig"),
+ ("dab-detr", "DabDetrConfig"),
+ ("dac", "DacConfig"),
+ ("data2vec-audio", "Data2VecAudioConfig"),
+ ("data2vec-text", "Data2VecTextConfig"),
+ ("data2vec-vision", "Data2VecVisionConfig"),
+ ("dbrx", "DbrxConfig"),
+ ("deberta", "DebertaConfig"),
+ ("deberta-v2", "DebertaV2Config"),
+ ("decision_transformer", "DecisionTransformerConfig"),
+ ("deepseek_v2", "DeepseekV2Config"),
+ ("deepseek_v3", "DeepseekV3Config"),
+ ("deepseek_vl", "DeepseekVLConfig"),
+ ("deepseek_vl_hybrid", "DeepseekVLHybridConfig"),
+ ("deformable_detr", "DeformableDetrConfig"),
+ ("deit", "DeiTConfig"),
+ ("depth_anything", "DepthAnythingConfig"),
+ ("depth_pro", "DepthProConfig"),
+ ("detr", "DetrConfig"),
+ ("dia", "DiaConfig"),
+ ("diffllama", "DiffLlamaConfig"),
+ ("dinat", "DinatConfig"),
+ ("dinov2", "Dinov2Config"),
+ ("dinov2_with_registers", "Dinov2WithRegistersConfig"),
+ ("dinov3_convnext", "DINOv3ConvNextConfig"),
+ ("dinov3_vit", "DINOv3ViTConfig"),
+ ("distilbert", "DistilBertConfig"),
+ ("doge", "DogeConfig"),
+ ("donut-swin", "DonutSwinConfig"),
+ ("dots1", "Dots1Config"),
+ ("dpr", "DPRConfig"),
+ ("dpt", "DPTConfig"),
+ ("edgetam", "EdgeTamConfig"),
+ ("edgetam_video", "EdgeTamVideoConfig"),
+ ("edgetam_vision_model", "EdgeTamVisionConfig"),
+ ("efficientloftr", "EfficientLoFTRConfig"),
+ ("efficientnet", "EfficientNetConfig"),
+ ("electra", "ElectraConfig"),
+ ("emu3", "Emu3Config"),
+ ("encodec", "EncodecConfig"),
+ ("encoder-decoder", "EncoderDecoderConfig"),
+ ("eomt", "EomtConfig"),
+ ("eomt_dinov3", "EomtDinov3Config"),
+ ("ernie", "ErnieConfig"),
+ ("ernie4_5", "Ernie4_5Config"),
+ ("ernie4_5_moe", "Ernie4_5_MoeConfig"),
+ ("ernie4_5_vl_moe", "Ernie4_5_VLMoeConfig"),
+ ("esm", "EsmConfig"),
+ ("eurobert", "EuroBertConfig"),
+ ("evolla", "EvollaConfig"),
+ ("exaone4", "Exaone4Config"),
+ ("exaone_moe", "ExaoneMoeConfig"),
+ ("falcon", "FalconConfig"),
+ ("falcon_h1", "FalconH1Config"),
+ ("falcon_mamba", "FalconMambaConfig"),
+ ("fast_vlm", "FastVlmConfig"),
+ ("fastspeech2_conformer", "FastSpeech2ConformerConfig"),
+ ("fastspeech2_conformer_with_hifigan", "FastSpeech2ConformerWithHifiGanConfig"),
+ ("flaubert", "FlaubertConfig"),
+ ("flava", "FlavaConfig"),
+ ("flex_olmo", "FlexOlmoConfig"),
+ ("florence2", "Florence2Config"),
+ ("fnet", "FNetConfig"),
+ ("focalnet", "FocalNetConfig"),
+ ("fsmt", "FSMTConfig"),
+ ("funnel", "FunnelConfig"),
+ ("fuyu", "FuyuConfig"),
+ ("gemma", "GemmaConfig"),
+ ("gemma2", "Gemma2Config"),
+ ("gemma3", "Gemma3Config"),
+ ("gemma3_text", "Gemma3TextConfig"),
+ ("gemma3n", "Gemma3nConfig"),
+ ("gemma3n_audio", "Gemma3nAudioConfig"),
+ ("gemma3n_text", "Gemma3nTextConfig"),
+ ("gemma3n_vision", "Gemma3nVisionConfig"),
+ ("gemma4", "Gemma4Config"),
+ ("gemma4_audio", "Gemma4AudioConfig"),
+ ("gemma4_text", "Gemma4TextConfig"),
+ ("gemma4_vision", "Gemma4VisionConfig"),
+ ("git", "GitConfig"),
+ ("glm", "GlmConfig"),
+ ("glm4", "Glm4Config"),
+ ("glm46v", "Glm46VConfig"),
+ ("glm4_moe", "Glm4MoeConfig"),
+ ("glm4_moe_lite", "Glm4MoeLiteConfig"),
+ ("glm4v", "Glm4vConfig"),
+ ("glm4v_moe", "Glm4vMoeConfig"),
+ ("glm4v_moe_text", "Glm4vMoeTextConfig"),
+ ("glm4v_moe_vision", "Glm4vMoeVisionConfig"),
+ ("glm4v_text", "Glm4vTextConfig"),
+ ("glm4v_vision", "Glm4vVisionConfig"),
+ ("glm_image", "GlmImageConfig"),
+ ("glm_image_text", "GlmImageTextConfig"),
+ ("glm_image_vision", "GlmImageVisionConfig"),
+ ("glm_image_vqmodel", "GlmImageVQVAEConfig"),
+ ("glm_moe_dsa", "GlmMoeDsaConfig"),
+ ("glm_ocr", "GlmOcrConfig"),
+ ("glm_ocr_text", "GlmOcrTextConfig"),
+ ("glm_ocr_vision", "GlmOcrVisionConfig"),
+ ("glmasr", "GlmAsrConfig"),
+ ("glmasr_encoder", "GlmAsrEncoderConfig"),
+ ("glpn", "GLPNConfig"),
+ ("got_ocr2", "GotOcr2Config"),
+ ("gpt-sw3", "GPT2Config"),
+ ("gpt2", "GPT2Config"),
+ ("gpt_bigcode", "GPTBigCodeConfig"),
+ ("gpt_neo", "GPTNeoConfig"),
+ ("gpt_neox", "GPTNeoXConfig"),
+ ("gpt_neox_japanese", "GPTNeoXJapaneseConfig"),
+ ("gpt_oss", "GptOssConfig"),
+ ("gptj", "GPTJConfig"),
+ ("granite", "GraniteConfig"),
+ ("granite_speech", "GraniteSpeechConfig"),
+ ("granitemoe", "GraniteMoeConfig"),
+ ("granitemoehybrid", "GraniteMoeHybridConfig"),
+ ("granitemoeshared", "GraniteMoeSharedConfig"),
+ ("granitevision", "LlavaNextConfig"),
+ ("grounding-dino", "GroundingDinoConfig"),
+ ("groupvit", "GroupViTConfig"),
+ ("helium", "HeliumConfig"),
+ ("hgnet_v2", "HGNetV2Config"),
+ ("hiera", "HieraConfig"),
+ ("higgs_audio_v2", "HiggsAudioV2Config"),
+ ("higgs_audio_v2_tokenizer", "HiggsAudioV2TokenizerConfig"),
+ ("hubert", "HubertConfig"),
+ ("hunyuan_v1_dense", "HunYuanDenseV1Config"),
+ ("hunyuan_v1_moe", "HunYuanMoEV1Config"),
+ ("ibert", "IBertConfig"),
+ ("idefics", "IdeficsConfig"),
+ ("idefics2", "Idefics2Config"),
+ ("idefics3", "Idefics3Config"),
+ ("idefics3_vision", "Idefics3VisionConfig"),
+ ("ijepa", "IJepaConfig"),
+ ("imagegpt", "ImageGPTConfig"),
+ ("informer", "InformerConfig"),
+ ("instructblip", "InstructBlipConfig"),
+ ("instructblipvideo", "InstructBlipVideoConfig"),
+ ("internvl", "InternVLConfig"),
+ ("internvl_vision", "InternVLVisionConfig"),
+ ("jais2", "Jais2Config"),
+ ("jamba", "JambaConfig"),
+ ("janus", "JanusConfig"),
+ ("jetmoe", "JetMoeConfig"),
+ ("jina_embeddings_v3", "JinaEmbeddingsV3Config"),
+ ("kosmos-2", "Kosmos2Config"),
+ ("kosmos-2.5", "Kosmos2_5Config"),
+ ("kyutai_speech_to_text", "KyutaiSpeechToTextConfig"),
+ ("lasr_ctc", "LasrCTCConfig"),
+ ("lasr_encoder", "LasrEncoderConfig"),
+ ("layoutlm", "LayoutLMConfig"),
+ ("layoutlmv2", "LayoutLMv2Config"),
+ ("layoutlmv3", "LayoutLMv3Config"),
+ ("layoutxlm", "LayoutXLMConfig"),
+ ("led", "LEDConfig"),
+ ("levit", "LevitConfig"),
+ ("lfm2", "Lfm2Config"),
+ ("lfm2_moe", "Lfm2MoeConfig"),
+ ("lfm2_vl", "Lfm2VlConfig"),
+ ("lightglue", "LightGlueConfig"),
+ ("lighton_ocr", "LightOnOcrConfig"),
+ ("lilt", "LiltConfig"),
+ ("llama", "LlamaConfig"),
+ ("llama4", "Llama4Config"),
+ ("llama4_text", "Llama4TextConfig"),
+ ("llava", "LlavaConfig"),
+ ("llava_next", "LlavaNextConfig"),
+ ("llava_next_video", "LlavaNextVideoConfig"),
+ ("llava_onevision", "LlavaOnevisionConfig"),
+ ("longcat_flash", "LongcatFlashConfig"),
+ ("longformer", "LongformerConfig"),
+ ("longt5", "LongT5Config"),
+ ("luke", "LukeConfig"),
+ ("lw_detr", "LwDetrConfig"),
+ ("lw_detr_vit", "LwDetrViTConfig"),
+ ("lxmert", "LxmertConfig"),
+ ("m2m_100", "M2M100Config"),
+ ("mamba", "MambaConfig"),
+ ("mamba2", "Mamba2Config"),
+ ("marian", "MarianConfig"),
+ ("markuplm", "MarkupLMConfig"),
+ ("mask2former", "Mask2FormerConfig"),
+ ("maskformer", "MaskFormerConfig"),
+ ("maskformer-swin", "MaskFormerSwinConfig"),
+ ("mbart", "MBartConfig"),
+ ("megatron-bert", "MegatronBertConfig"),
+ ("metaclip_2", "MetaClip2Config"),
+ ("mgp-str", "MgpstrConfig"),
+ ("mimi", "MimiConfig"),
+ ("minimax", "MiniMaxConfig"),
+ ("minimax_m2", "MiniMaxM2Config"),
+ ("ministral", "MinistralConfig"),
+ ("ministral3", "Ministral3Config"),
+ ("mistral", "MistralConfig"),
+ ("mistral3", "Mistral3Config"),
+ ("mistral4", "Mistral4Config"),
+ ("mixtral", "MixtralConfig"),
+ ("mlcd", "MLCDVisionConfig"), # Keep this to make some original hub repositories (from `DeepGlint-AI`) works
+ ("mlcd_vision_model", "MLCDVisionConfig"),
+ ("mllama", "MllamaConfig"),
+ ("mm-grounding-dino", "MMGroundingDinoConfig"),
+ ("mobilebert", "MobileBertConfig"),
+ ("mobilenet_v1", "MobileNetV1Config"),
+ ("mobilenet_v2", "MobileNetV2Config"),
+ ("mobilevit", "MobileViTConfig"),
+ ("mobilevitv2", "MobileViTV2Config"),
+ ("modernbert", "ModernBertConfig"),
+ ("modernbert-decoder", "ModernBertDecoderConfig"),
+ ("modernvbert", "ModernVBertConfig"),
+ ("moonshine", "MoonshineConfig"),
+ ("moonshine_streaming", "MoonshineStreamingConfig"),
+ ("moonshine_streaming_encoder", "MoonshineStreamingEncoderConfig"),
+ ("moshi", "MoshiConfig"),
+ ("mpnet", "MPNetConfig"),
+ ("mpt", "MptConfig"),
+ ("mra", "MraConfig"),
+ ("mt5", "MT5Config"),
+ ("musicflamingo", "MusicFlamingoConfig"),
+ ("musicflamingo_encoder", "AudioFlamingo3EncoderConfig"),
+ ("musicgen", "MusicgenConfig"),
+ ("musicgen_melody", "MusicgenMelodyConfig"),
+ ("mvp", "MvpConfig"),
+ ("nanochat", "NanoChatConfig"),
+ ("nemotron", "NemotronConfig"),
+ ("nemotron_h", "NemotronHConfig"),
+ ("nllb-moe", "NllbMoeConfig"),
+ ("nomic_bert", "NomicBertConfig"),
+ ("nougat", "VisionEncoderDecoderConfig"),
+ ("nystromformer", "NystromformerConfig"),
+ ("olmo", "OlmoConfig"),
+ ("olmo2", "Olmo2Config"),
+ ("olmo3", "Olmo3Config"),
+ ("olmo_hybrid", "OlmoHybridConfig"),
+ ("olmoe", "OlmoeConfig"),
+ ("omdet-turbo", "OmDetTurboConfig"),
+ ("oneformer", "OneFormerConfig"),
+ ("openai-gpt", "OpenAIGPTConfig"),
+ ("opt", "OPTConfig"),
+ ("ovis2", "Ovis2Config"),
+ ("owlv2", "Owlv2Config"),
+ ("owlvit", "OwlViTConfig"),
+ ("paddleocr_vl", "PaddleOCRVLConfig"),
+ ("paligemma", "PaliGemmaConfig"),
+ ("parakeet_ctc", "ParakeetCTCConfig"),
+ ("parakeet_encoder", "ParakeetEncoderConfig"),
+ ("patchtsmixer", "PatchTSMixerConfig"),
+ ("patchtst", "PatchTSTConfig"),
+ ("pe_audio", "PeAudioConfig"),
+ ("pe_audio_encoder", "PeAudioEncoderConfig"),
+ ("pe_audio_video", "PeAudioVideoConfig"),
+ ("pe_audio_video_encoder", "PeAudioVideoEncoderConfig"),
+ ("pe_video", "PeVideoConfig"),
+ ("pe_video_encoder", "PeVideoEncoderConfig"),
+ ("pegasus", "PegasusConfig"),
+ ("pegasus_x", "PegasusXConfig"),
+ ("perceiver", "PerceiverConfig"),
+ ("perception_lm", "PerceptionLMConfig"),
+ ("persimmon", "PersimmonConfig"),
+ ("phi", "PhiConfig"),
+ ("phi3", "Phi3Config"),
+ ("phi4_multimodal", "Phi4MultimodalConfig"),
+ ("phimoe", "PhimoeConfig"),
+ ("pi0", "PI0Config"),
+ ("pix2struct", "Pix2StructConfig"),
+ ("pixio", "PixioConfig"),
+ ("pixtral", "PixtralVisionConfig"),
+ ("plbart", "PLBartConfig"),
+ ("poolformer", "PoolFormerConfig"),
+ ("pop2piano", "Pop2PianoConfig"),
+ ("pp_chart2table", "PPChart2TableConfig"),
+ ("pp_doclayout_v2", "PPDocLayoutV2Config"),
+ ("pp_doclayout_v3", "PPDocLayoutV3Config"),
+ ("pp_lcnet", "PPLCNetConfig"),
+ ("pp_lcnet_v3", "PPLCNetV3Config"),
+ ("pp_ocrv5_mobile_det", "PPOCRV5MobileDetConfig"),
+ ("pp_ocrv5_mobile_rec", "PPOCRV5MobileRecConfig"),
+ ("pp_ocrv5_server_det", "PPOCRV5ServerDetConfig"),
+ ("pp_ocrv5_server_rec", "PPOCRV5ServerRecConfig"),
+ ("prompt_depth_anything", "PromptDepthAnythingConfig"),
+ ("prophetnet", "ProphetNetConfig"),
+ ("pvt", "PvtConfig"),
+ ("pvt_v2", "PvtV2Config"),
+ ("qwen2", "Qwen2Config"),
+ ("qwen2_5_omni", "Qwen2_5OmniConfig"),
+ ("qwen2_5_vl", "Qwen2_5_VLConfig"),
+ ("qwen2_5_vl_text", "Qwen2_5_VLTextConfig"),
+ ("qwen2_audio", "Qwen2AudioConfig"),
+ ("qwen2_audio_encoder", "Qwen2AudioEncoderConfig"),
+ ("qwen2_moe", "Qwen2MoeConfig"),
+ ("qwen2_vl", "Qwen2VLConfig"),
+ ("qwen2_vl_text", "Qwen2VLTextConfig"),
+ ("qwen3", "Qwen3Config"),
+ ("qwen3_5", "Qwen3_5Config"),
+ ("qwen3_5_moe", "Qwen3_5MoeConfig"),
+ ("qwen3_5_moe_text", "Qwen3_5MoeTextConfig"),
+ ("qwen3_5_text", "Qwen3_5TextConfig"),
+ ("qwen3_moe", "Qwen3MoeConfig"),
+ ("qwen3_next", "Qwen3NextConfig"),
+ ("qwen3_omni_moe", "Qwen3OmniMoeConfig"),
+ ("qwen3_vl", "Qwen3VLConfig"),
+ ("qwen3_vl_moe", "Qwen3VLMoeConfig"),
+ ("qwen3_vl_moe_text", "Qwen3VLMoeTextConfig"),
+ ("qwen3_vl_text", "Qwen3VLTextConfig"),
+ ("rag", "RagConfig"),
+ ("recurrent_gemma", "RecurrentGemmaConfig"),
+ ("reformer", "ReformerConfig"),
+ ("regnet", "RegNetConfig"),
+ ("rembert", "RemBertConfig"),
+ ("resnet", "ResNetConfig"),
+ ("roberta", "RobertaConfig"),
+ ("roberta-prelayernorm", "RobertaPreLayerNormConfig"),
+ ("roc_bert", "RoCBertConfig"),
+ ("roformer", "RoFormerConfig"),
+ ("rt_detr", "RTDetrConfig"),
+ ("rt_detr_resnet", "RTDetrResNetConfig"),
+ ("rt_detr_v2", "RTDetrV2Config"),
+ ("rwkv", "RwkvConfig"),
+ ("sam", "SamConfig"),
+ ("sam2", "Sam2Config"),
+ ("sam2_hiera_det_model", "Sam2HieraDetConfig"),
+ ("sam2_video", "Sam2VideoConfig"),
+ ("sam2_vision_model", "Sam2VisionConfig"),
+ ("sam3", "Sam3Config"),
+ ("sam3_tracker", "Sam3TrackerConfig"),
+ ("sam3_tracker_video", "Sam3TrackerVideoConfig"),
+ ("sam3_video", "Sam3VideoConfig"),
+ ("sam3_vision_model", "Sam3VisionConfig"),
+ ("sam3_vit_model", "Sam3ViTConfig"),
+ ("sam_hq", "SamHQConfig"),
+ ("sam_hq_vision_model", "SamHQVisionConfig"),
+ ("sam_vision_model", "SamVisionConfig"),
+ ("seamless_m4t", "SeamlessM4TConfig"),
+ ("seamless_m4t_v2", "SeamlessM4Tv2Config"),
+ ("seed_oss", "SeedOssConfig"),
+ ("segformer", "SegformerConfig"),
+ ("seggpt", "SegGptConfig"),
+ ("sew", "SEWConfig"),
+ ("sew-d", "SEWDConfig"),
+ ("shieldgemma2", "ShieldGemma2Config"),
+ ("siglip", "SiglipConfig"),
+ ("siglip2", "Siglip2Config"),
+ ("siglip2_vision_model", "Siglip2VisionConfig"),
+ ("siglip_vision_model", "SiglipVisionConfig"),
+ ("slanext", "SLANeXtConfig"),
+ ("smollm3", "SmolLM3Config"),
+ ("smolvlm", "SmolVLMConfig"),
+ ("smolvlm_vision", "SmolVLMVisionConfig"),
+ ("solar_open", "SolarOpenConfig"),
+ ("speech-encoder-decoder", "SpeechEncoderDecoderConfig"),
+ ("speech_to_text", "Speech2TextConfig"),
+ ("speecht5", "SpeechT5Config"),
+ ("splinter", "SplinterConfig"),
+ ("squeezebert", "SqueezeBertConfig"),
+ ("stablelm", "StableLmConfig"),
+ ("starcoder2", "Starcoder2Config"),
+ ("superglue", "SuperGlueConfig"),
+ ("superpoint", "SuperPointConfig"),
+ ("swiftformer", "SwiftFormerConfig"),
+ ("swin", "SwinConfig"),
+ ("swin2sr", "Swin2SRConfig"),
+ ("swinv2", "Swinv2Config"),
+ ("switch_transformers", "SwitchTransformersConfig"),
+ ("t5", "T5Config"),
+ ("t5gemma", "T5GemmaConfig"),
+ ("t5gemma2", "T5Gemma2Config"),
+ ("t5gemma2_encoder", "T5Gemma2EncoderConfig"),
+ ("table-transformer", "TableTransformerConfig"),
+ ("tapas", "TapasConfig"),
+ ("textnet", "TextNetConfig"),
+ ("time_series_transformer", "TimeSeriesTransformerConfig"),
+ ("timesfm", "TimesFmConfig"),
+ ("timesfm2_5", "TimesFm2_5Config"),
+ ("timesformer", "TimesformerConfig"),
+ ("timm_backbone", "TimmBackboneConfig"),
+ ("timm_wrapper", "TimmWrapperConfig"),
+ ("trocr", "TrOCRConfig"),
+ ("tvp", "TvpConfig"),
+ ("udop", "UdopConfig"),
+ ("umt5", "UMT5Config"),
+ ("unispeech", "UniSpeechConfig"),
+ ("unispeech-sat", "UniSpeechSatConfig"),
+ ("univnet", "UnivNetConfig"),
+ ("upernet", "UperNetConfig"),
+ ("uvdoc", "UVDocConfig"),
+ ("uvdoc_backbone", "UVDocBackboneConfig"),
+ ("vaultgemma", "VaultGemmaConfig"),
+ ("vibevoice_acoustic_tokenizer", "VibeVoiceAcousticTokenizerConfig"),
+ ("vibevoice_acoustic_tokenizer_decoder", "VibeVoiceAcousticTokenizerDecoderConfig"),
+ ("vibevoice_acoustic_tokenizer_encoder", "VibeVoiceAcousticTokenizerEncoderConfig"),
+ ("vibevoice_asr", "VibeVoiceAsrConfig"),
+ ("video_llama_3", "VideoLlama3Config"),
+ ("video_llama_3_vision", "VideoLlama3VisionConfig"),
+ ("video_llava", "VideoLlavaConfig"),
+ ("videomae", "VideoMAEConfig"),
+ ("videomt", "VideomtConfig"),
+ ("vilt", "ViltConfig"),
+ ("vipllava", "VipLlavaConfig"),
+ ("vision-encoder-decoder", "VisionEncoderDecoderConfig"),
+ ("vision-text-dual-encoder", "VisionTextDualEncoderConfig"),
+ ("visual_bert", "VisualBertConfig"),
+ ("vit", "ViTConfig"),
+ ("vit_mae", "ViTMAEConfig"),
+ ("vit_msn", "ViTMSNConfig"),
+ ("vitdet", "VitDetConfig"),
+ ("vitmatte", "VitMatteConfig"),
+ ("vitpose", "VitPoseConfig"),
+ ("vitpose_backbone", "VitPoseBackboneConfig"),
+ ("vits", "VitsConfig"),
+ ("vivit", "VivitConfig"),
+ ("vjepa2", "VJEPA2Config"),
+ ("voxtral", "VoxtralConfig"),
+ ("voxtral_encoder", "VoxtralEncoderConfig"),
+ ("voxtral_realtime", "VoxtralRealtimeConfig"),
+ ("voxtral_realtime_encoder", "VoxtralRealtimeEncoderConfig"),
+ ("voxtral_realtime_text", "VoxtralRealtimeTextConfig"),
+ ("wav2vec2", "Wav2Vec2Config"),
+ ("wav2vec2-bert", "Wav2Vec2BertConfig"),
+ ("wav2vec2-conformer", "Wav2Vec2ConformerConfig"),
+ ("wavlm", "WavLMConfig"),
+ ("whisper", "WhisperConfig"),
+ ("xclip", "XCLIPConfig"),
+ ("xcodec", "XcodecConfig"),
+ ("xglm", "XGLMConfig"),
+ ("xlm", "XLMConfig"),
+ ("xlm-roberta", "XLMRobertaConfig"),
+ ("xlm-roberta-xl", "XLMRobertaXLConfig"),
+ ("xlnet", "XLNetConfig"),
+ ("xlstm", "xLSTMConfig"),
+ ("xmod", "XmodConfig"),
+ ("yolos", "YolosConfig"),
+ ("yoso", "YosoConfig"),
+ ("youtu", "YoutuConfig"),
+ ("zamba", "ZambaConfig"),
+ ("zamba2", "Zamba2Config"),
+ ("zoedepth", "ZoeDepthConfig"),
+ ]
+)
+
+
+MODEL_NAMES_MAPPING = OrderedDict[str, str](
+ [
+ # Add full (and cased) model names here
+ ("afmoe", "AFMoE"),
+ ("aimv2", "AIMv2"),
+ ("aimv2_vision_model", "Aimv2VisionModel"),
+ ("albert", "ALBERT"),
+ ("align", "ALIGN"),
+ ("altclip", "AltCLIP"),
+ ("apertus", "Apertus"),
+ ("arcee", "Arcee"),
+ ("aria", "Aria"),
+ ("aria_text", "AriaText"),
+ ("audio-spectrogram-transformer", "Audio Spectrogram Transformer"),
+ ("audioflamingo3", "AudioFlamingo3"),
+ ("audioflamingo3_encoder", "AudioFlamingo3Encoder"),
+ ("autoformer", "Autoformer"),
+ ("aya_vision", "AyaVision"),
+ ("bamba", "Bamba"),
+ ("bark", "Bark"),
+ ("bart", "BART"),
+ ("barthez", "BARThez"),
+ ("bartpho", "BARTpho"),
+ ("beit", "BEiT"),
+ ("bert", "BERT"),
+ ("bert-generation", "Bert Generation"),
+ ("bert-japanese", "BertJapanese"),
+ ("bertweet", "BERTweet"),
+ ("big_bird", "BigBird"),
+ ("bigbird_pegasus", "BigBird-Pegasus"),
+ ("biogpt", "BioGpt"),
+ ("bit", "BiT"),
+ ("bitnet", "BitNet"),
+ ("blenderbot", "Blenderbot"),
+ ("blenderbot-small", "BlenderbotSmall"),
+ ("blip", "BLIP"),
+ ("blip-2", "BLIP-2"),
+ ("blip_2_qformer", "BLIP-2 QFormer"),
+ ("bloom", "BLOOM"),
+ ("blt", "Blt"),
+ ("bridgetower", "BridgeTower"),
+ ("bros", "BROS"),
+ ("byt5", "ByT5"),
+ ("camembert", "CamemBERT"),
+ ("canine", "CANINE"),
+ ("chameleon", "Chameleon"),
+ ("chinese_clip", "Chinese-CLIP"),
+ ("chinese_clip_vision_model", "ChineseCLIPVisionModel"),
+ ("chmv2", "CHMv2"),
+ ("clap", "CLAP"),
+ ("clip", "CLIP"),
+ ("clip_text_model", "CLIPTextModel"),
+ ("clip_vision_model", "CLIPVisionModel"),
+ ("clipseg", "CLIPSeg"),
+ ("clvp", "CLVP"),
+ ("code_llama", "CodeLlama"),
+ ("codegen", "CodeGen"),
+ ("cohere", "Cohere"),
+ ("cohere2", "Cohere2"),
+ ("cohere2_vision", "Cohere2Vision"),
+ ("cohere_asr", "CohereASR"),
+ ("colmodernvbert", "ColModernVBert"),
+ ("colpali", "ColPali"),
+ ("colqwen2", "ColQwen2"),
+ ("conditional_detr", "Conditional DETR"),
+ ("convbert", "ConvBERT"),
+ ("convnext", "ConvNeXT"),
+ ("convnextv2", "ConvNeXTV2"),
+ ("cpm", "CPM"),
+ ("cpmant", "CPM-Ant"),
+ ("csm", "CSM"),
+ ("ctrl", "CTRL"),
+ ("cvt", "CvT"),
+ ("cwm", "Code World Model (CWM)"),
+ ("d_fine", "D-FINE"),
+ ("dab-detr", "DAB-DETR"),
+ ("dac", "DAC"),
+ ("data2vec-audio", "Data2VecAudio"),
+ ("data2vec-text", "Data2VecText"),
+ ("data2vec-vision", "Data2VecVision"),
+ ("dbrx", "DBRX"),
+ ("deberta", "DeBERTa"),
+ ("deberta-v2", "DeBERTa-v2"),
+ ("decision_transformer", "Decision Transformer"),
+ ("deepseek_v2", "DeepSeek-V2"),
+ ("deepseek_v3", "DeepSeek-V3"),
+ ("deepseek_vl", "DeepseekVL"),
+ ("deepseek_vl_hybrid", "DeepseekVLHybrid"),
+ ("deformable_detr", "Deformable DETR"),
+ ("deit", "DeiT"),
+ ("deplot", "DePlot"),
+ ("depth_anything", "Depth Anything"),
+ ("depth_anything_v2", "Depth Anything V2"),
+ ("depth_pro", "DepthPro"),
+ ("detr", "DETR"),
+ ("dia", "Dia"),
+ ("dialogpt", "DialoGPT"),
+ ("diffllama", "DiffLlama"),
+ ("dinat", "DiNAT"),
+ ("dinov2", "DINOv2"),
+ ("dinov2_with_registers", "DINOv2 with Registers"),
+ ("dinov3_convnext", "DINOv3 ConvNext"),
+ ("dinov3_vit", "DINOv3 ViT"),
+ ("distilbert", "DistilBERT"),
+ ("dit", "DiT"),
+ ("doge", "Doge"),
+ ("donut-swin", "DonutSwin"),
+ ("dots1", "dots1"),
+ ("dpr", "DPR"),
+ ("dpt", "DPT"),
+ ("edgetam", "EdgeTAM"),
+ ("edgetam_video", "EdgeTamVideo"),
+ ("edgetam_vision_model", "EdgeTamVisionModel"),
+ ("efficientloftr", "EfficientLoFTR"),
+ ("efficientnet", "EfficientNet"),
+ ("electra", "ELECTRA"),
+ ("emu3", "Emu3"),
+ ("encodec", "EnCodec"),
+ ("encoder-decoder", "Encoder decoder"),
+ ("eomt", "EoMT"),
+ ("eomt_dinov3", "EoMT-DINOv3"),
+ ("ernie", "ERNIE"),
+ ("ernie4_5", "Ernie4_5"),
+ ("ernie4_5_moe", "Ernie4_5_MoE"),
+ ("ernie4_5_vl_moe", "Ernie4_5_VLMoE"),
+ ("esm", "ESM"),
+ ("eurobert", "EuroBERT"),
+ ("evolla", "Evolla"),
+ ("exaone4", "EXAONE-4.0"),
+ ("exaone_moe", "EXAONE-MoE"),
+ ("falcon", "Falcon"),
+ ("falcon3", "Falcon3"),
+ ("falcon_h1", "FalconH1"),
+ ("falcon_mamba", "FalconMamba"),
+ ("fast_vlm", "FastVlm"),
+ ("fastspeech2_conformer", "FastSpeech2Conformer"),
+ ("fastspeech2_conformer_with_hifigan", "FastSpeech2ConformerWithHifiGan"),
+ ("flan-t5", "FLAN-T5"),
+ ("flan-ul2", "FLAN-UL2"),
+ ("flaubert", "FlauBERT"),
+ ("flava", "FLAVA"),
+ ("flex_olmo", "FlexOlmo"),
+ ("florence2", "Florence2"),
+ ("fnet", "FNet"),
+ ("focalnet", "FocalNet"),
+ ("fsmt", "FairSeq Machine-Translation"),
+ ("funnel", "Funnel Transformer"),
+ ("fuyu", "Fuyu"),
+ ("gemma", "Gemma"),
+ ("gemma2", "Gemma2"),
+ ("gemma3", "Gemma3ForConditionalGeneration"),
+ ("gemma3_text", "Gemma3ForCausalLM"),
+ ("gemma3n", "Gemma3nForConditionalGeneration"),
+ ("gemma3n_audio", "Gemma3nAudioEncoder"),
+ ("gemma3n_text", "Gemma3nForCausalLM"),
+ ("gemma3n_vision", "TimmWrapperModel"),
+ ("gemma4", "Gemma4ForConditionalGeneration"),
+ ("gemma4_audio", "Gemma4AudioModel"),
+ ("gemma4_text", "Gemma4ForCausalLM"),
+ ("gemma4_vision", "Gemma4VisionModel"),
+ ("git", "GIT"),
+ ("glm", "GLM"),
+ ("glm4", "GLM4"),
+ ("glm46v", "Glm46V"),
+ ("glm4_moe", "Glm4MoE"),
+ ("glm4_moe_lite", "Glm4MoELite"),
+ ("glm4v", "GLM4V"),
+ ("glm4v_moe", "GLM4VMOE"),
+ ("glm4v_moe_text", "GLM4VMOE"),
+ ("glm4v_moe_vision", "Glm4vMoeVisionModel"),
+ ("glm4v_text", "GLM4V"),
+ ("glm4v_vision", "Glm4vVisionModel"),
+ ("glm_image", "GlmImage"),
+ ("glm_image_text", "GlmImageText"),
+ ("glm_image_vision", "GlmImageVisionModel"),
+ ("glm_image_vqmodel", "GlmImageVQVAE"),
+ ("glm_moe_dsa", "GlmMoeDsa"),
+ ("glm_ocr", "Glmocr"),
+ ("glm_ocr_text", "GlmOcrText"),
+ ("glm_ocr_vision", "GlmOcrVisionModel"),
+ ("glmasr", "GLM-ASR"),
+ ("glmasr_encoder", "GLM-ASR Encoder"),
+ ("glpn", "GLPN"),
+ ("got_ocr2", "GOT-OCR2"),
+ ("gpt-sw3", "GPT-Sw3"),
+ ("gpt2", "OpenAI GPT-2"),
+ ("gpt_bigcode", "GPTBigCode"),
+ ("gpt_neo", "GPT Neo"),
+ ("gpt_neox", "GPT NeoX"),
+ ("gpt_neox_japanese", "GPT NeoX Japanese"),
+ ("gpt_oss", "GptOss"),
+ ("gptj", "GPT-J"),
+ ("granite", "Granite"),
+ ("granite_speech", "GraniteSpeech"),
+ ("granitemoe", "GraniteMoeMoe"),
+ ("granitemoehybrid", "GraniteMoeHybrid"),
+ ("granitemoeshared", "GraniteMoeSharedMoe"),
+ ("granitevision", "LLaVA-NeXT"),
+ ("grounding-dino", "Grounding DINO"),
+ ("groupvit", "GroupViT"),
+ ("helium", "Helium"),
+ ("herbert", "HerBERT"),
+ ("hgnet_v2", "HGNet-V2"),
+ ("hiera", "Hiera"),
+ ("higgs_audio_v2", "HiggsAudioV2"),
+ ("higgs_audio_v2_tokenizer", "HiggsAudioV2Tokenizer"),
+ ("hubert", "Hubert"),
+ ("hunyuan_v1_dense", "HunYuanDenseV1"),
+ ("hunyuan_v1_moe", "HunYuanMoeV1"),
+ ("ibert", "I-BERT"),
+ ("idefics", "IDEFICS"),
+ ("idefics2", "Idefics2"),
+ ("idefics3", "Idefics3"),
+ ("idefics3_vision", "Idefics3VisionTransformer"),
+ ("ijepa", "I-JEPA"),
+ ("imagegpt", "ImageGPT"),
+ ("informer", "Informer"),
+ ("instructblip", "InstructBLIP"),
+ ("instructblipvideo", "InstructBlipVideo"),
+ ("internvl", "InternVL"),
+ ("internvl_vision", "InternVLVision"),
+ ("jais2", "Jais2"),
+ ("jamba", "Jamba"),
+ ("janus", "Janus"),
+ ("jetmoe", "JetMoe"),
+ ("jina_embeddings_v3", "JinaEmbeddingsV3"),
+ ("kosmos-2", "KOSMOS-2"),
+ ("kosmos-2.5", "KOSMOS-2.5"),
+ ("kyutai_speech_to_text", "KyutaiSpeechToText"),
+ ("lasr", "Lasr"),
+ ("lasr_ctc", "Lasr"),
+ ("lasr_encoder", "LasrEncoder"),
+ ("layoutlm", "LayoutLM"),
+ ("layoutlmv2", "LayoutLMv2"),
+ ("layoutlmv3", "LayoutLMv3"),
+ ("layoutxlm", "LayoutXLM"),
+ ("led", "LED"),
+ ("levit", "LeViT"),
+ ("lfm2", "Lfm2"),
+ ("lfm2_moe", "Lfm2Moe"),
+ ("lfm2_vl", "Lfm2Vl"),
+ ("lightglue", "LightGlue"),
+ ("lighton_ocr", "LightOnOcr"),
+ ("lilt", "LiLT"),
+ ("llama", "LLaMA"),
+ ("llama2", "Llama2"),
+ ("llama3", "Llama3"),
+ ("llama4", "Llama4"),
+ ("llama4_text", "Llama4ForCausalLM"),
+ ("llava", "LLaVa"),
+ ("llava_next", "LLaVA-NeXT"),
+ ("llava_next_video", "LLaVa-NeXT-Video"),
+ ("llava_onevision", "LLaVA-Onevision"),
+ ("longcat_flash", "LongCatFlash"),
+ ("longformer", "Longformer"),
+ ("longt5", "LongT5"),
+ ("luke", "LUKE"),
+ ("lw_detr", "LwDetr"),
+ ("lw_detr_vit", "LwDetrVit"),
+ ("lxmert", "LXMERT"),
+ ("m2m_100", "M2M100"),
+ ("madlad-400", "MADLAD-400"),
+ ("mamba", "Mamba"),
+ ("mamba2", "mamba2"),
+ ("marian", "Marian"),
+ ("markuplm", "MarkupLM"),
+ ("mask2former", "Mask2Former"),
+ ("maskformer", "MaskFormer"),
+ ("maskformer-swin", "MaskFormerSwin"),
+ ("matcha", "MatCha"),
+ ("mbart", "mBART"),
+ ("mbart50", "mBART-50"),
+ ("megatron-bert", "Megatron-BERT"),
+ ("megatron_gpt2", "Megatron-GPT2"),
+ ("metaclip_2", "MetaCLIP 2"),
+ ("mgp-str", "MGP-STR"),
+ ("mimi", "Mimi"),
+ ("minimax", "MiniMax"),
+ ("minimax_m2", "MiniMax-M2"),
+ ("ministral", "Ministral"),
+ ("ministral3", "Ministral3"),
+ ("mistral", "Mistral"),
+ ("mistral3", "Mistral3"),
+ ("mistral4", "Mistral4"),
+ ("mixtral", "Mixtral"),
+ ("mlcd", "MLCD"), # Keep this to make some original hub repositories (from `DeepGlint-AI`) works
+ ("mlcd_vision_model", "MLCD"),
+ ("mllama", "Mllama"),
+ ("mluke", "mLUKE"),
+ ("mm-grounding-dino", "MM Grounding DINO"),
+ ("mms", "MMS"),
+ ("mobilebert", "MobileBERT"),
+ ("mobilenet_v1", "MobileNetV1"),
+ ("mobilenet_v2", "MobileNetV2"),
+ ("mobilevit", "MobileViT"),
+ ("mobilevitv2", "MobileViTV2"),
+ ("modernbert", "ModernBERT"),
+ ("modernbert-decoder", "ModernBertDecoder"),
+ ("modernvbert", "ModernVBert"),
+ ("moonshine", "Moonshine"),
+ ("moonshine_streaming", "MoonshineStreaming"),
+ ("moonshine_streaming_encoder", "MoonshineStreamingEncoder"),
+ ("moshi", "Moshi"),
+ ("mpnet", "MPNet"),
+ ("mpt", "MPT"),
+ ("mra", "MRA"),
+ ("mt5", "MT5"),
+ ("musicflamingo", "MusicFlamingo"),
+ ("musicflamingo_encoder", "AudioFlamingo3Encoder"),
+ ("musicgen", "MusicGen"),
+ ("musicgen_melody", "MusicGen Melody"),
+ ("mvp", "MVP"),
+ ("myt5", "myt5"),
+ ("nanochat", "NanoChat"),
+ ("nemotron", "Nemotron"),
+ ("nemotron_h", "NemotronH"),
+ ("nllb", "NLLB"),
+ ("nllb-moe", "NLLB-MOE"),
+ ("nomic_bert", "NomicBERT"),
+ ("nougat", "Nougat"),
+ ("nystromformer", "Nyströmformer"),
+ ("olmo", "OLMo"),
+ ("olmo2", "OLMo2"),
+ ("olmo3", "Olmo3"),
+ ("olmo_hybrid", "OlmoHybrid"),
+ ("olmoe", "OLMoE"),
+ ("omdet-turbo", "OmDet-Turbo"),
+ ("oneformer", "OneFormer"),
+ ("openai-gpt", "OpenAI GPT"),
+ ("opt", "OPT"),
+ ("ovis2", "Ovis2"),
+ ("owlv2", "OWLv2"),
+ ("owlvit", "OWL-ViT"),
+ ("paddleocr_vl", "PaddleOCRVL"),
+ ("paligemma", "PaliGemma"),
+ ("parakeet", "Parakeet"),
+ ("parakeet_ctc", "Parakeet"),
+ ("parakeet_encoder", "ParakeetEncoder"),
+ ("patchtsmixer", "PatchTSMixer"),
+ ("patchtst", "PatchTST"),
+ ("pe_audio", "PeAudio"),
+ ("pe_audio_encoder", "PeAudioEncoder"),
+ ("pe_audio_video", "PeAudioVideo"),
+ ("pe_audio_video_encoder", "PeAudioVideoEncoder"),
+ ("pe_video", "PeVideo"),
+ ("pe_video_encoder", "PeVideoEncoder"),
+ ("pegasus", "Pegasus"),
+ ("pegasus_x", "PEGASUS-X"),
+ ("perceiver", "Perceiver"),
+ ("perception_lm", "PerceptionLM"),
+ ("persimmon", "Persimmon"),
+ ("phi", "Phi"),
+ ("phi3", "Phi3"),
+ ("phi4_multimodal", "Phi4Multimodal"),
+ ("phimoe", "Phimoe"),
+ ("phobert", "PhoBERT"),
+ ("pi0", "PI0"),
+ ("pix2struct", "Pix2Struct"),
+ ("pixio", "Pixio"),
+ ("pixtral", "Pixtral"),
+ ("plbart", "PLBart"),
+ ("poolformer", "PoolFormer"),
+ ("pop2piano", "Pop2Piano"),
+ ("pp_chart2table", "PPChart2Table"),
+ ("pp_doclayout_v2", "PPDocLayoutV2"),
+ ("pp_doclayout_v3", "PPDocLayoutV3"),
+ ("pp_lcnet", "PPLCNet"),
+ ("pp_lcnet_v3", "PPLCNetV3"),
+ ("pp_ocrv5_mobile_det", "PPOCRV5MobileDet"),
+ ("pp_ocrv5_mobile_rec", "PPOCRV5MobileRec"),
+ ("pp_ocrv5_server_det", "PPOCRV5ServerDet"),
+ ("pp_ocrv5_server_rec", "PPOCRV5ServerRec"),
+ ("prompt_depth_anything", "PromptDepthAnything"),
+ ("prophetnet", "ProphetNet"),
+ ("pvt", "PVT"),
+ ("pvt_v2", "PVTv2"),
+ ("qwen2", "Qwen2"),
+ ("qwen2_5_omni", "Qwen2_5Omni"),
+ ("qwen2_5_vl", "Qwen2_5_VL"),
+ ("qwen2_5_vl_text", "Qwen2_5_VL"),
+ ("qwen2_audio", "Qwen2Audio"),
+ ("qwen2_audio_encoder", "Qwen2AudioEncoder"),
+ ("qwen2_moe", "Qwen2MoE"),
+ ("qwen2_vl", "Qwen2VL"),
+ ("qwen2_vl_text", "Qwen2VL"),
+ ("qwen3", "Qwen3"),
+ ("qwen3_5", "Qwen3_5"),
+ ("qwen3_5_moe", "Qwen3_5Moe"),
+ ("qwen3_5_moe_text", "Qwen3_5MoeText"),
+ ("qwen3_5_text", "Qwen3_5Text"),
+ ("qwen3_moe", "Qwen3MoE"),
+ ("qwen3_next", "Qwen3Next"),
+ ("qwen3_omni_moe", "Qwen3OmniMoE"),
+ ("qwen3_vl", "Qwen3VL"),
+ ("qwen3_vl_moe", "Qwen3VLMoe"),
+ ("qwen3_vl_moe_text", "Qwen3VLMoe"),
+ ("qwen3_vl_text", "Qwen3VL"),
+ ("rag", "RAG"),
+ ("recurrent_gemma", "RecurrentGemma"),
+ ("reformer", "Reformer"),
+ ("regnet", "RegNet"),
+ ("rembert", "RemBERT"),
+ ("resnet", "ResNet"),
+ ("roberta", "RoBERTa"),
+ ("roberta-prelayernorm", "RoBERTa-PreLayerNorm"),
+ ("roc_bert", "RoCBert"),
+ ("roformer", "RoFormer"),
+ ("rt_detr", "RT-DETR"),
+ ("rt_detr_resnet", "RT-DETR-ResNet"),
+ ("rt_detr_v2", "RT-DETRv2"),
+ ("rwkv", "RWKV"),
+ ("sam", "SAM"),
+ ("sam2", "SAM2"),
+ ("sam2_hiera_det_model", "Sam2HieraDetModel"),
+ ("sam2_video", "Sam2VideoModel"),
+ ("sam2_vision_model", "Sam2VisionModel"),
+ ("sam3", "SAM3"),
+ ("sam3_tracker", "Sam3Tracker"),
+ ("sam3_tracker_video", "Sam3TrackerVideo"),
+ ("sam3_video", "Sam3VideoModel"),
+ ("sam3_vision_model", "Sam3VisionModel"),
+ ("sam3_vit_model", "Sam3ViTModel"),
+ ("sam_hq", "SAM-HQ"),
+ ("sam_hq_vision_model", "SamHQVisionModel"),
+ ("sam_vision_model", "SamVisionModel"),
+ ("seamless_m4t", "SeamlessM4T"),
+ ("seamless_m4t_v2", "SeamlessM4Tv2"),
+ ("seed_oss", "SeedOss"),
+ ("segformer", "SegFormer"),
+ ("seggpt", "SegGPT"),
+ ("sew", "SEW"),
+ ("sew-d", "SEW-D"),
+ ("shieldgemma2", "Shieldgemma2"),
+ ("siglip", "SigLIP"),
+ ("siglip2", "SigLIP2"),
+ ("siglip2_vision_model", "Siglip2VisionModel"),
+ ("siglip_vision_model", "SiglipVisionModel"),
+ ("slanext", "SLANeXt"),
+ ("smollm3", "SmolLM3"),
+ ("smolvlm", "SmolVLM"),
+ ("smolvlm_vision", "SmolVLMVisionTransformer"),
+ ("solar_open", "SolarOpen"),
+ ("speech-encoder-decoder", "Speech Encoder decoder"),
+ ("speech_to_text", "Speech2Text"),
+ ("speecht5", "SpeechT5"),
+ ("splinter", "Splinter"),
+ ("squeezebert", "SqueezeBERT"),
+ ("stablelm", "StableLm"),
+ ("starcoder2", "Starcoder2"),
+ ("superglue", "SuperGlue"),
+ ("superpoint", "SuperPoint"),
+ ("swiftformer", "SwiftFormer"),
+ ("swin", "Swin Transformer"),
+ ("swin2sr", "Swin2SR"),
+ ("swinv2", "Swin Transformer V2"),
+ ("switch_transformers", "SwitchTransformers"),
+ ("t5", "T5"),
+ ("t5gemma", "T5Gemma"),
+ ("t5gemma2", "T5Gemma2"),
+ ("t5gemma2_encoder", "T5Gemma2Encoder"),
+ ("t5v1.1", "T5v1.1"),
+ ("table-transformer", "Table Transformer"),
+ ("tapas", "TAPAS"),
+ ("textnet", "TextNet"),
+ ("time_series_transformer", "Time Series Transformer"),
+ ("timesfm", "TimesFm"),
+ ("timesfm2_5", "TimesFm2p5"),
+ ("timesformer", "TimeSformer"),
+ ("timm_backbone", "TimmBackbone"),
+ ("timm_wrapper", "TimmWrapperModel"),
+ ("trocr", "TrOCR"),
+ ("tvp", "TVP"),
+ ("udop", "UDOP"),
+ ("ul2", "UL2"),
+ ("umt5", "UMT5"),
+ ("unispeech", "UniSpeech"),
+ ("unispeech-sat", "UniSpeechSat"),
+ ("univnet", "UnivNet"),
+ ("upernet", "UPerNet"),
+ ("uvdoc", "UVDoc"),
+ ("uvdoc_backbone", "UVDocBackbone"),
+ ("vaultgemma", "VaultGemma"),
+ ("vibevoice_acoustic_tokenizer", "VibeVoiceAcousticTokenizer"),
+ ("vibevoice_acoustic_tokenizer_decoder", "VibeVoiceAcousticTokenizerDecoderConfig"),
+ ("vibevoice_acoustic_tokenizer_encoder", "VibeVoiceAcousticTokenizerEncoderConfig"),
+ ("vibevoice_asr", "VibeVoiceAsr"),
+ ("video_llama_3", "VideoLlama3"),
+ ("video_llama_3_vision", "VideoLlama3Vision"),
+ ("video_llava", "VideoLlava"),
+ ("videomae", "VideoMAE"),
+ ("videomt", "VidEoMT"),
+ ("vilt", "ViLT"),
+ ("vipllava", "VipLlava"),
+ ("vision-encoder-decoder", "Vision Encoder decoder"),
+ ("vision-text-dual-encoder", "VisionTextDualEncoder"),
+ ("visual_bert", "VisualBERT"),
+ ("vit", "ViT"),
+ ("vit_mae", "ViTMAE"),
+ ("vit_msn", "ViTMSN"),
+ ("vitdet", "VitDet"),
+ ("vitmatte", "ViTMatte"),
+ ("vitpose", "ViTPose"),
+ ("vitpose_backbone", "ViTPoseBackbone"),
+ ("vits", "VITS"),
+ ("vivit", "ViViT"),
+ ("vjepa2", "VJEPA2Model"),
+ ("voxtral", "Voxtral"),
+ ("voxtral_encoder", "Voxtral Encoder"),
+ ("voxtral_realtime", "VoxtralRealtime"),
+ ("voxtral_realtime_encoder", "VoxtralRealtime Encoder"),
+ ("voxtral_realtime_text", "VoxtralRealtime Text Model"),
+ ("wav2vec2", "Wav2Vec2"),
+ ("wav2vec2-bert", "Wav2Vec2-BERT"),
+ ("wav2vec2-conformer", "Wav2Vec2-Conformer"),
+ ("wav2vec2_phoneme", "Wav2Vec2Phoneme"),
+ ("wavlm", "WavLM"),
+ ("whisper", "Whisper"),
+ ("xclip", "X-CLIP"),
+ ("xcodec", "X-CODEC"),
+ ("xglm", "XGLM"),
+ ("xlm", "XLM"),
+ ("xlm-roberta", "XLM-RoBERTa"),
+ ("xlm-roberta-xl", "XLM-RoBERTa-XL"),
+ ("xlm-v", "XLM-V"),
+ ("xlnet", "XLNet"),
+ ("xls_r", "XLS-R"),
+ ("xlsr_wav2vec2", "XLSR-Wav2Vec2"),
+ ("xlstm", "xLSTM"),
+ ("xmod", "X-MOD"),
+ ("yolos", "YOLOS"),
+ ("yoso", "YOSO"),
+ ("youtu", "Youtu"),
+ ("zamba", "Zamba"),
+ ("zamba2", "Zamba2"),
+ ("zoedepth", "ZoeDepth"),
+ ]
+)
+
+# This is tied to the processing `-` -> `_` in `model_type_to_module_name`. For example, instead of putting
+# `transfo-xl` (as in `CONFIG_MAPPING_NAMES`), we should use `transfo_xl`.
+DEPRECATED_MODELS = []
+
+SPECIAL_MODEL_TYPE_TO_MODULE_NAME = OrderedDict[str, str](
+ [
+ ("audioflamingo3_encoder", "audioflamingo3"),
+ ("musicflamingo_encoder", "musicflamingo"),
+ ("openai-gpt", "openai"),
+ ("blip-2", "blip_2"),
+ ("data2vec-audio", "data2vec"),
+ ("data2vec-text", "data2vec"),
+ ("data2vec-vision", "data2vec"),
+ ("donut-swin", "donut"),
+ ("kosmos-2", "kosmos2"),
+ ("kosmos-2.5", "kosmos2_5"),
+ ("mlcd_vision_model", "mlcd"),
+ ("omdet-turbo", "omdet_turbo"),
+ ("maskformer-swin", "maskformer"),
+ ("xclip", "x_clip"),
+ ("clip_vision_model", "clip"),
+ ("qwen2_audio_encoder", "qwen2_audio"),
+ ("voxtral_encoder", "voxtral"),
+ ("voxtral_realtime_encoder", "voxtral_realtime"),
+ ("voxtral_realtime_text", "voxtral_realtime"),
+ ("clip_text_model", "clip"),
+ ("aria_text", "aria"),
+ ("gemma3_text", "gemma3"),
+ ("gemma3n_audio", "gemma3n"),
+ ("gemma3n_text", "gemma3n"),
+ ("gemma3n_vision", "gemma3n"),
+ ("gemma4_audio", "gemma4"),
+ ("gemma4_text", "gemma4"),
+ ("gemma4_vision", "gemma4"),
+ ("glm4v_vision", "glm4v"),
+ ("glm4v_moe_vision", "glm4v_moe"),
+ ("glm4v_text", "glm4v"),
+ ("glm4v_moe_text", "glm4v_moe"),
+ ("glm_image_vision", "glm_image"),
+ ("glm_image_vqmodel", "glm_image"),
+ ("glm_image_text", "glm_image"),
+ ("glm_ocr_vision", "glm_ocr"),
+ ("glm_ocr_vqmodel", "glm_ocr"),
+ ("glm_ocr_text", "glm_ocr"),
+ ("glmasr_encoder", "glmasr"),
+ ("grounding-dino", "grounding_dino"),
+ ("moonshine_streaming_encoder", "moonshine_streaming"),
+ ("mm-grounding-dino", "mm_grounding_dino"),
+ ("idefics3_vision", "idefics3"),
+ ("mgp-str", "mgp_str"),
+ ("siglip_vision_model", "siglip"),
+ ("siglip2_vision_model", "siglip2"),
+ ("aimv2_vision_model", "aimv2"),
+ ("smolvlm_vision", "smolvlm"),
+ ("chinese_clip_vision_model", "chinese_clip"),
+ ("rt_detr_resnet", "rt_detr"),
+ ("granitevision", "llava_next"),
+ ("internvl_vision", "internvl"),
+ ("qwen2_5_vl_text", "qwen2_5_vl"),
+ ("qwen2_vl_text", "qwen2_vl"),
+ ("qwen3_vl_text", "qwen3_vl"),
+ ("qwen3_vl_moe_text", "qwen3_vl_moe"),
+ ("qwen3_5_text", "qwen3_5"),
+ ("qwen3_5_moe_text", "qwen3_5_moe"),
+ ("sam_vision_model", "sam"),
+ ("sam2_vision_model", "sam2"),
+ ("sam2_hiera_det_model", "sam2"),
+ ("sam3_vit_model", "sam3"),
+ ("sam3_vision_model", "sam3"),
+ ("edgetam_vision_model", "edgetam"),
+ ("sam_hq_vision_model", "sam_hq"),
+ ("t5gemma2_encoder", "t5gemma2"),
+ ("llama4_text", "llama4"),
+ ("blip_2_qformer", "blip_2"),
+ ("fastspeech2_conformer_with_hifigan", "fastspeech2_conformer"),
+ ("perception_encoder", "perception_lm"),
+ ("pe_audio_encoder", "pe_audio"),
+ ("pe_video_encoder", "pe_video"),
+ ("pe_audio_video_encoder", "pe_audio_video"),
+ ("video_llama_3_vision", "video_llama_3"),
+ ("parakeet_encoder", "parakeet"),
+ ("lw_detr_vit", "lw_detr"),
+ ("parakeet_ctc", "parakeet"),
+ ("lasr_encoder", "lasr"),
+ ("lasr_ctc", "lasr"),
+ ("wav2vec2-bert", "wav2vec2_bert"),
+ ("vibevoice_acoustic_tokenizer_encoder", "vibevoice_acoustic_tokenizer"),
+ ("vibevoice_acoustic_tokenizer_decoder", "vibevoice_acoustic_tokenizer"),
+ ("uvdoc_backbone", "uvdoc"),
+ ]
+)
+
+
+def model_type_to_module_name(key) -> str:
+ """Converts a config key to the corresponding module."""
+ # Special treatment
+ if key in SPECIAL_MODEL_TYPE_TO_MODULE_NAME:
+ key = SPECIAL_MODEL_TYPE_TO_MODULE_NAME[key]
+
+ if key in DEPRECATED_MODELS:
+ key = f"deprecated.{key}"
+ return key
+
+ key = key.replace("-", "_")
+ if key in DEPRECATED_MODELS:
+ key = f"deprecated.{key}"
+
+ return key
+
+
+def config_class_to_model_type(config) -> str | None:
+ """Converts a config class name to the corresponding model type"""
+ for key, cls in CONFIG_MAPPING_NAMES.items():
+ if cls == config:
+ return key
+ # if key not found check in extra content
+ for key, cls in CONFIG_MAPPING._extra_content.items():
+ if cls.__name__ == config:
+ return key
+ return None
+
+
+class _LazyConfigMapping(OrderedDict[str, type[PreTrainedConfig]]):
+ """
+ A dictionary that lazily load its values when they are requested.
+ """
+
+ def __init__(self, mapping) -> None:
+ self._mapping = mapping
+ self._extra_content = {}
+ self._modules = {}
+
+ def __getitem__(self, key: str) -> type[PreTrainedConfig]:
+ if key in self._extra_content:
+ return self._extra_content[key]
+ if key not in self._mapping:
+ raise KeyError(key)
+ value = self._mapping[key]
+ module_name = model_type_to_module_name(key)
+ if module_name not in self._modules:
+ self._modules[module_name] = importlib.import_module(f".{module_name}", "transformers.models")
+ if hasattr(self._modules[module_name], value):
+ return getattr(self._modules[module_name], value)
+
+ # Some of the mappings have entries model_type -> config of another model type. In that case we try to grab the
+ # object at the top level.
+ transformers_module = importlib.import_module("transformers")
+ return getattr(transformers_module, value)
+
+ def keys(self) -> list[str]:
+ return list(self._mapping.keys()) + list(self._extra_content.keys())
+
+ def values(self) -> list[type[PreTrainedConfig]]:
+ return [self[k] for k in self._mapping] + list(self._extra_content.values())
+
+ def items(self) -> list[tuple[str, type[PreTrainedConfig]]]:
+ return [(k, self[k]) for k in self._mapping] + list(self._extra_content.items())
+
+ def __iter__(self) -> Iterator[str]:
+ return iter(list(self._mapping.keys()) + list(self._extra_content.keys()))
+
+ def __contains__(self, item: object) -> bool:
+ return item in self._mapping or item in self._extra_content
+
+ def register(self, key: str, value: type[PreTrainedConfig], exist_ok=False) -> None:
+ """
+ Register a new configuration in this mapping.
+ """
+ if key in self._mapping and not exist_ok:
+ raise ValueError(f"'{key}' is already used by a Transformers config, pick another name.")
+ self._extra_content[key] = value
+
+
+CONFIG_MAPPING = _LazyConfigMapping(CONFIG_MAPPING_NAMES)
+
+
+class _LazyLoadAllMappings(OrderedDict[str, str]):
+ """
+ A mapping that will load all pairs of key values at the first access (either by indexing, requestions keys, values,
+ etc.)
+
+ Args:
+ mapping: The mapping to load.
+ """
+
+ def __init__(self, mapping):
+ self._mapping = mapping
+ self._initialized = False
+ self._data = {}
+
+ def _initialize(self):
+ if self._initialized:
+ return
+
+ for model_type, map_name in self._mapping.items():
+ module_name = model_type_to_module_name(model_type)
+ module = importlib.import_module(f".{module_name}", "transformers.models")
+ mapping = getattr(module, map_name)
+ self._data.update(mapping)
+
+ self._initialized = True
+
+ def __getitem__(self, key):
+ self._initialize()
+ return self._data[key]
+
+ def keys(self) -> KeysView[str]:
+ self._initialize()
+ return self._data.keys()
+
+ def values(self) -> ValuesView[str]:
+ self._initialize()
+ return self._data.values()
+
+ def items(self) -> KeysView[str]:
+ self._initialize()
+ return self._data.keys()
+
+ def __iter__(self) -> Iterator[str]:
+ self._initialize()
+ return iter(self._data)
+
+ def __contains__(self, item: object) -> bool:
+ self._initialize()
+ return item in self._data
+
+
+def _get_class_name(model_class: str | list[str]):
+ if isinstance(model_class, (list, tuple)):
+ return " or ".join([f"[`{c}`]" for c in model_class if c is not None])
+ return f"[`{model_class}`]"
+
+
+def _list_model_options(indent, config_to_class=None, use_model_types=True):
+ if config_to_class is None and not use_model_types:
+ raise ValueError("Using `use_model_types=False` requires a `config_to_class` dictionary.")
+ if use_model_types:
+ if config_to_class is None:
+ model_type_to_name = {model_type: f"[`{config}`]" for model_type, config in CONFIG_MAPPING_NAMES.items()}
+ else:
+ model_type_to_name = {
+ model_type: _get_class_name(model_class)
+ for model_type, model_class in config_to_class.items()
+ if model_type in MODEL_NAMES_MAPPING
+ }
+ lines = [
+ f"{indent}- **{model_type}** -- {model_type_to_name[model_type]} ({MODEL_NAMES_MAPPING[model_type]} model)"
+ for model_type in sorted(model_type_to_name.keys())
+ ]
+ else:
+ config_to_name = {
+ CONFIG_MAPPING_NAMES[config]: _get_class_name(clas)
+ for config, clas in config_to_class.items()
+ if config in CONFIG_MAPPING_NAMES
+ }
+ config_to_model_name = {
+ config: MODEL_NAMES_MAPPING[model_type] for model_type, config in CONFIG_MAPPING_NAMES.items()
+ }
+ lines = [
+ f"{indent}- [`{config_name}`] configuration class:"
+ f" {config_to_name[config_name]} ({config_to_model_name[config_name]} model)"
+ for config_name in sorted(config_to_name.keys())
+ ]
+ return "\n".join(lines)
+
+
+def replace_list_option_in_docstrings(
+ config_to_class=None, use_model_types: bool = True
+) -> Callable[[_CallableT], _CallableT]:
+ def docstring_decorator(fn):
+ docstrings = fn.__doc__
+ if docstrings is None:
+ # Example: -OO
+ return fn
+ lines = docstrings.split("\n")
+ i = 0
+ while i < len(lines) and re.search(r"^(\s*)List options\s*$", lines[i]) is None:
+ i += 1
+ if i < len(lines):
+ indent = re.search(r"^(\s*)List options\s*$", lines[i]).groups()[0]
+ if use_model_types:
+ indent = f"{indent} "
+ lines[i] = _list_model_options(indent, config_to_class=config_to_class, use_model_types=use_model_types)
+ docstrings = "\n".join(lines)
+ else:
+ raise ValueError(
+ f"The function {fn} should have an empty 'List options' in its docstring as placeholder, current"
+ f" docstring is:\n{docstrings}"
+ )
+ fn.__doc__ = docstrings
+ return fn
+
+ return docstring_decorator
+
+
+class AutoConfig:
+ r"""
+ This is a generic configuration class that will be instantiated as one of the configuration classes of the library
+ when created with the [`~AutoConfig.from_pretrained`] class method.
+
+ This class cannot be instantiated directly using `__init__()` (throws an error).
+ """
+
+ def __init__(self) -> None:
+ raise OSError(
+ "AutoConfig is designed to be instantiated "
+ "using the `AutoConfig.from_pretrained(pretrained_model_name_or_path)` method."
+ )
+
+ @classmethod
+ def for_model(cls, model_type: str, *args, **kwargs) -> PreTrainedConfig:
+ if model_type in CONFIG_MAPPING:
+ config_class = CONFIG_MAPPING[model_type]
+ return config_class(*args, **kwargs)
+ raise ValueError(
+ f"Unrecognized model identifier: {model_type}. Should contain one of {', '.join(CONFIG_MAPPING.keys())}"
+ )
+
+ @classmethod
+ @replace_list_option_in_docstrings()
+ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike[str], **kwargs):
+ r"""
+ Instantiate one of the configuration classes of the library from a pretrained model configuration.
+
+ The configuration class to instantiate is selected based on the `model_type` property of the config object that
+ is loaded, or when it's missing, by falling back to using pattern matching on `pretrained_model_name_or_path`:
+
+ List options
+
+ Args:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ Can be either:
+
+ - A string, the *model id* of a pretrained model configuration hosted inside a model repo on
+ huggingface.co.
+ - A path to a *directory* containing a configuration file saved using the
+ [`~PreTrainedConfig.save_pretrained`] method, or the [`~PreTrainedModel.save_pretrained`] method,
+ e.g., `./my_model_directory/`.
+ - a path to a saved configuration JSON *file*, e.g.,
+ `./my_model_directory/configuration.json`.
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the
+ standard cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force the (re-)download the model weights and configuration files and override the
+ cached versions if they exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ return_unused_kwargs (`bool`, *optional*, defaults to `False`):
+ If `False`, then this function returns just the final configuration object.
+
+ If `True`, then this functions returns a `Tuple(config, unused_kwargs)` where *unused_kwargs* is a
+ dictionary consisting of the key/value pairs whose keys are not configuration attributes: i.e., the
+ part of `kwargs` which has not been used to update `config` and is otherwise ignored.
+ trust_remote_code (`bool`, *optional*, defaults to `False`):
+ Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
+ should only be set to `True` for repositories you trust and in which you have read the code, as it will
+ execute code present on the Hub on your local machine.
+ kwargs(additional keyword arguments, *optional*):
+ The values in kwargs of any keys which are configuration attributes will be used to override the loaded
+ values. Behavior concerning key/value pairs whose keys are *not* configuration attributes is controlled
+ by the `return_unused_kwargs` keyword parameter.
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoConfig
+
+ >>> # Download configuration from huggingface.co and cache.
+ >>> config = AutoConfig.from_pretrained("google-bert/bert-base-uncased")
+
+ >>> # Download configuration from huggingface.co (user-uploaded) and cache.
+ >>> config = AutoConfig.from_pretrained("dbmdz/bert-base-german-cased")
+
+ >>> # If configuration file is in a directory (e.g., was saved using *save_pretrained('./test/saved_model/')*).
+ >>> config = AutoConfig.from_pretrained("./test/bert_saved_model/")
+
+ >>> # Load a specific configuration file.
+ >>> config = AutoConfig.from_pretrained("./test/bert_saved_model/my_configuration.json")
+
+ >>> # Change some config attributes when loading a pretrained config.
+ >>> config = AutoConfig.from_pretrained("google-bert/bert-base-uncased", output_attentions=True, foo=False)
+ >>> config.output_attentions
+ True
+
+ >>> config, unused_kwargs = AutoConfig.from_pretrained(
+ ... "google-bert/bert-base-uncased", output_attentions=True, foo=False, return_unused_kwargs=True
+ ... )
+ >>> config.output_attentions
+ True
+
+ >>> unused_kwargs
+ {'foo': False}
+ ```
+ """
+ kwargs["_from_auto"] = True
+ kwargs["name_or_path"] = pretrained_model_name_or_path
+ trust_remote_code = kwargs.pop("trust_remote_code", None)
+ code_revision = kwargs.pop("code_revision", None)
+
+ config_dict, unused_kwargs = PreTrainedConfig.get_config_dict(pretrained_model_name_or_path, **kwargs)
+ has_remote_code = "auto_map" in config_dict and "AutoConfig" in config_dict["auto_map"]
+ has_local_code = "model_type" in config_dict and config_dict["model_type"] in CONFIG_MAPPING
+ explicit_local_code = has_local_code and not CONFIG_MAPPING[config_dict["model_type"]].__module__.startswith(
+ "transformers."
+ )
+ if has_remote_code:
+ class_ref = config_dict["auto_map"]["AutoConfig"]
+ if "--" in class_ref:
+ upstream_repo = class_ref.split("--")[0]
+ else:
+ upstream_repo = None
+ trust_remote_code = resolve_trust_remote_code(
+ trust_remote_code, pretrained_model_name_or_path, has_local_code, has_remote_code, upstream_repo
+ )
+
+ if has_remote_code and trust_remote_code and not explicit_local_code:
+ config_class = get_class_from_dynamic_module(
+ class_ref, pretrained_model_name_or_path, code_revision=code_revision, **kwargs
+ )
+ config_class.register_for_auto_class()
+ return config_class.from_pretrained(pretrained_model_name_or_path, **kwargs)
+ elif "model_type" in config_dict:
+ # Apply heuristic: if model_type is mistral but layer_types is present, treat as ministral
+ if config_dict["model_type"] == "mistral" and "layer_types" in config_dict:
+ logger.info(
+ "Detected mistral model with layer_types, treating as ministral for alternating attention compatibility. "
+ )
+ config_dict["model_type"] = "ministral"
+
+ try:
+ config_class = CONFIG_MAPPING[config_dict["model_type"]]
+ except KeyError:
+ raise ValueError(
+ f"The checkpoint you are trying to load has model type `{config_dict['model_type']}` "
+ "but Transformers does not recognize this architecture. This could be because of an "
+ "issue with the checkpoint, or because your version of Transformers is out of date.\n\n"
+ "You can update Transformers with the command `pip install --upgrade transformers`. If this "
+ "does not work, and the checkpoint is very new, then there may not be a release version "
+ "that supports this model yet. In this case, you can get the most up-to-date code by installing "
+ "Transformers from source with the command "
+ "`pip install git+https://github.com/huggingface/transformers.git`"
+ )
+ return config_class.from_dict(config_dict, **unused_kwargs)
+
+ raise ValueError(
+ f"Unrecognized model in {pretrained_model_name_or_path}. "
+ f"Should have a `model_type` key in its {CONFIG_NAME}."
+ )
+
+ @staticmethod
+ def register(model_type, config, exist_ok=False) -> None:
+ """
+ Register a new configuration for this class.
+
+ Args:
+ model_type (`str`): The model type like "bert" or "gpt".
+ config ([`PreTrainedConfig`]): The config to register.
+ """
+ if issubclass(config, PreTrainedConfig) and config.model_type != model_type:
+ raise ValueError(
+ "The config you are passing has a `model_type` attribute that is not consistent with the model type "
+ f"you passed (config has {config.model_type} and you passed {model_type}. Fix one of those so they "
+ "match!"
+ )
+ CONFIG_MAPPING.register(model_type, config, exist_ok=exist_ok)
+
+
+__all__ = ["CONFIG_MAPPING", "MODEL_NAMES_MAPPING", "AutoConfig"]
diff --git a/third_party/transformers/src/transformers/models/auto/feature_extraction_auto.py b/third_party/transformers/src/transformers/models/auto/feature_extraction_auto.py
new file mode 100644
index 0000000000000000000000000000000000000000..111c56efb43625205c00ffbb3ed502e7bfccc388
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/auto/feature_extraction_auto.py
@@ -0,0 +1,386 @@
+# Copyright 2021 The HuggingFace Inc. team.
+#
+# 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.
+"""AutoFeatureExtractor class."""
+
+import importlib
+import os
+from collections import OrderedDict
+
+# Build the list of all feature extractors
+from ...configuration_utils import PreTrainedConfig
+from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
+from ...feature_extraction_utils import FeatureExtractionMixin
+from ...utils import CONFIG_NAME, FEATURE_EXTRACTOR_NAME, PROCESSOR_NAME, cached_file, logging, safe_load_json_file
+from .auto_factory import _LazyAutoMapping
+from .configuration_auto import (
+ CONFIG_MAPPING_NAMES,
+ AutoConfig,
+ model_type_to_module_name,
+ replace_list_option_in_docstrings,
+)
+
+
+logger = logging.get_logger(__name__)
+
+FEATURE_EXTRACTOR_MAPPING_NAMES = OrderedDict(
+ [
+ ("audio-spectrogram-transformer", "ASTFeatureExtractor"),
+ ("audioflamingo3", "WhisperFeatureExtractor"),
+ ("clap", "ClapFeatureExtractor"),
+ ("clvp", "ClvpFeatureExtractor"),
+ ("cohere_asr", "CohereAsrFeatureExtractor"),
+ ("csm", "EncodecFeatureExtractor"),
+ ("dac", "DacFeatureExtractor"),
+ ("data2vec-audio", "Wav2Vec2FeatureExtractor"),
+ ("dia", "DiaFeatureExtractor"),
+ ("encodec", "EncodecFeatureExtractor"),
+ ("gemma3n", "Gemma3nAudioFeatureExtractor"),
+ ("gemma4", "Gemma4AudioFeatureExtractor"),
+ ("glmasr", "WhisperFeatureExtractor"),
+ ("granite_speech", "GraniteSpeechFeatureExtractor"),
+ ("higgs_audio_v2_tokenizer", "DacFeatureExtractor"),
+ ("hubert", "Wav2Vec2FeatureExtractor"),
+ ("kyutai_speech_to_text", "KyutaiSpeechToTextFeatureExtractor"),
+ ("lasr_ctc", "LasrFeatureExtractor"),
+ ("lasr_encoder", "LasrFeatureExtractor"),
+ ("markuplm", "MarkupLMFeatureExtractor"),
+ ("mimi", "EncodecFeatureExtractor"),
+ ("moonshine", "Wav2Vec2FeatureExtractor"),
+ ("moshi", "EncodecFeatureExtractor"),
+ ("musicgen", "EncodecFeatureExtractor"),
+ ("musicgen_melody", "MusicgenMelodyFeatureExtractor"),
+ ("parakeet_ctc", "ParakeetFeatureExtractor"),
+ ("parakeet_encoder", "ParakeetFeatureExtractor"),
+ ("pe_audio", "PeAudioFeatureExtractor"),
+ ("pe_audio_video", "PeAudioFeatureExtractor"),
+ ("phi4_multimodal", "Phi4MultimodalFeatureExtractor"),
+ ("pop2piano", "Pop2PianoFeatureExtractor"),
+ ("qwen2_5_omni", "WhisperFeatureExtractor"),
+ ("qwen2_audio", "WhisperFeatureExtractor"),
+ ("qwen3_omni_moe", "WhisperFeatureExtractor"),
+ ("seamless_m4t", "SeamlessM4TFeatureExtractor"),
+ ("seamless_m4t_v2", "SeamlessM4TFeatureExtractor"),
+ ("sew", "Wav2Vec2FeatureExtractor"),
+ ("sew-d", "Wav2Vec2FeatureExtractor"),
+ ("speech_to_text", "Speech2TextFeatureExtractor"),
+ ("speecht5", "SpeechT5FeatureExtractor"),
+ ("unispeech", "Wav2Vec2FeatureExtractor"),
+ ("unispeech-sat", "Wav2Vec2FeatureExtractor"),
+ ("univnet", "UnivNetFeatureExtractor"),
+ ("vibevoice_acoustic_tokenizer", "VibeVoiceAcousticTokenizerFeatureExtractor"),
+ ("vibevoice_asr", "VibeVoiceAcousticTokenizerFeatureExtractor"),
+ ("voxtral", "WhisperFeatureExtractor"),
+ ("voxtral_realtime", "VoxtralRealtimeFeatureExtractor"),
+ ("wav2vec2", "Wav2Vec2FeatureExtractor"),
+ ("wav2vec2-bert", "Wav2Vec2FeatureExtractor"),
+ ("wav2vec2-conformer", "Wav2Vec2FeatureExtractor"),
+ ("wavlm", "Wav2Vec2FeatureExtractor"),
+ ("whisper", "WhisperFeatureExtractor"),
+ ("xcodec", "DacFeatureExtractor"),
+ ]
+)
+
+FEATURE_EXTRACTOR_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, FEATURE_EXTRACTOR_MAPPING_NAMES)
+
+
+def feature_extractor_class_from_name(class_name: str):
+ for module_name, extractors in FEATURE_EXTRACTOR_MAPPING_NAMES.items():
+ if class_name in extractors:
+ module_name = model_type_to_module_name(module_name)
+
+ module = importlib.import_module(f".{module_name}", "transformers.models")
+ try:
+ return getattr(module, class_name)
+ except AttributeError:
+ continue
+
+ for extractor in FEATURE_EXTRACTOR_MAPPING._extra_content.values():
+ if getattr(extractor, "__name__", None) == class_name:
+ return extractor
+
+ # We did not find the class, but maybe it's because a dep is missing. In that case, the class will be in the main
+ # init and we return the proper dummy to get an appropriate error message.
+ main_module = importlib.import_module("transformers")
+ if hasattr(main_module, class_name):
+ return getattr(main_module, class_name)
+
+ return None
+
+
+def get_feature_extractor_config(
+ pretrained_model_name_or_path: str | os.PathLike,
+ cache_dir: str | os.PathLike | None = None,
+ force_download: bool = False,
+ proxies: dict[str, str] | None = None,
+ token: bool | str | None = None,
+ revision: str | None = None,
+ local_files_only: bool = False,
+ **kwargs,
+):
+ """
+ Loads the feature extractor configuration from a pretrained model feature extractor configuration.
+
+ Args:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained model configuration hosted inside a model repo on
+ huggingface.co.
+ - a path to a *directory* containing a configuration file saved using the
+ [`~FeatureExtractionMixin.save_pretrained`] method, e.g., `./my_model_directory/`.
+
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
+ cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the configuration files and override the cached versions if they
+ exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `hf auth login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ local_files_only (`bool`, *optional*, defaults to `False`):
+ If `True`, will only try to load the feature extractor configuration from local files.
+
+
+
+ Passing `token=True` is required when you want to use a private model.
+
+
+
+ Returns:
+ `Dict`: The configuration of the feature extractor.
+
+ Examples:
+
+ ```python
+ # Download configuration from huggingface.co and cache.
+ feature_extractor_config = get_feature_extractor_config("facebook/wav2vec2-base-960h")
+ # This model does not have a feature extractor config so the result will be an empty dict.
+ feature_extractor_config = get_feature_extractor_config("FacebookAI/xlm-roberta-base")
+
+ # Save a pretrained feature extractor locally and you can reload its config
+ from transformers import AutoFeatureExtractor
+
+ feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base-960h")
+ feature_extractor.save_pretrained("feature-extractor-test")
+ feature_extractor_config = get_feature_extractor_config("feature-extractor-test")
+ ```"""
+ # Load with a priority given to the nested processor config, if available in repo
+ resolved_processor_file = cached_file(
+ pretrained_model_name_or_path,
+ filename=PROCESSOR_NAME,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ token=token,
+ revision=revision,
+ local_files_only=local_files_only,
+ _raise_exceptions_for_gated_repo=False,
+ _raise_exceptions_for_missing_entries=False,
+ )
+ resolved_feature_extractor_file = cached_file(
+ pretrained_model_name_or_path,
+ filename=FEATURE_EXTRACTOR_NAME,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ token=token,
+ revision=revision,
+ local_files_only=local_files_only,
+ _raise_exceptions_for_gated_repo=False,
+ _raise_exceptions_for_missing_entries=False,
+ )
+
+ # An empty list if none of the possible files is found in the repo
+ if not resolved_feature_extractor_file and not resolved_processor_file:
+ logger.info("Could not locate the feature extractor configuration file.")
+ return {}
+
+ # Load feature_extractor dict. Priority goes as (nested config if found -> feature extractor config)
+ # We are downloading both configs because almost all models have a `processor_config.json` but
+ # not all of these are nested. We need to check if it was saved recently as nested or if it is legacy style
+ feature_extractor_dict = {}
+ if resolved_processor_file is not None:
+ processor_dict = safe_load_json_file(resolved_processor_file)
+ if "feature_extractor" in processor_dict:
+ feature_extractor_dict = processor_dict["feature_extractor"]
+
+ if resolved_feature_extractor_file is not None and feature_extractor_dict is None:
+ feature_extractor_dict = safe_load_json_file(resolved_feature_extractor_file)
+ return feature_extractor_dict
+
+
+class AutoFeatureExtractor:
+ r"""
+ This is a generic feature extractor class that will be instantiated as one of the feature extractor classes of the
+ library when created with the [`AutoFeatureExtractor.from_pretrained`] class method.
+
+ This class cannot be instantiated directly using `__init__()` (throws an error).
+ """
+
+ def __init__(self):
+ raise OSError(
+ "AutoFeatureExtractor is designed to be instantiated "
+ "using the `AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path)` method."
+ )
+
+ @classmethod
+ @replace_list_option_in_docstrings(FEATURE_EXTRACTOR_MAPPING_NAMES)
+ def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+ r"""
+ Instantiate one of the feature extractor classes of the library from a pretrained model vocabulary.
+
+ The feature extractor class to instantiate is selected based on the `model_type` property of the config object
+ (either passed as an argument or loaded from `pretrained_model_name_or_path` if possible), or when it's
+ missing, by falling back to using pattern matching on `pretrained_model_name_or_path`:
+
+ List options
+
+ Params:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained feature_extractor hosted inside a model repo on
+ huggingface.co.
+ - a path to a *directory* containing a feature extractor file saved using the
+ [`~feature_extraction_utils.FeatureExtractionMixin.save_pretrained`] method, e.g.,
+ `./my_model_directory/`.
+ - a path to a saved feature extractor JSON *file*, e.g.,
+ `./my_model_directory/preprocessor_config.json`.
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model feature extractor should be cached if the
+ standard cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the feature extractor files and override the cached versions
+ if they exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `hf auth login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ return_unused_kwargs (`bool`, *optional*, defaults to `False`):
+ If `False`, then this function returns just the final feature extractor object. If `True`, then this
+ functions returns a `Tuple(feature_extractor, unused_kwargs)` where *unused_kwargs* is a dictionary
+ consisting of the key/value pairs whose keys are not feature extractor attributes: i.e., the part of
+ `kwargs` which has not been used to update `feature_extractor` and is otherwise ignored.
+ trust_remote_code (`bool`, *optional*, defaults to `False`):
+ Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
+ should only be set to `True` for repositories you trust and in which you have read the code, as it will
+ execute code present on the Hub on your local machine.
+ kwargs (`dict[str, Any]`, *optional*):
+ The values in kwargs of any keys which are feature extractor attributes will be used to override the
+ loaded values. Behavior concerning key/value pairs whose keys are *not* feature extractor attributes is
+ controlled by the `return_unused_kwargs` keyword parameter.
+
+
+
+ Passing `token=True` is required when you want to use a private model.
+
+
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoFeatureExtractor
+
+ >>> # Download feature extractor from huggingface.co and cache.
+ >>> feature_extractor = AutoFeatureExtractor.from_pretrained("facebook/wav2vec2-base-960h")
+
+ >>> # If feature extractor files are in a directory (e.g. feature extractor was saved using *save_pretrained('./test/saved_model/')*)
+ >>> # feature_extractor = AutoFeatureExtractor.from_pretrained("./test/saved_model/")
+ ```"""
+ config = kwargs.pop("config", None)
+ trust_remote_code = kwargs.pop("trust_remote_code", None)
+ kwargs["_from_auto"] = True
+
+ config_dict, _ = FeatureExtractionMixin.get_feature_extractor_dict(pretrained_model_name_or_path, **kwargs)
+ feature_extractor_class = config_dict.get("feature_extractor_type", None)
+ feature_extractor_auto_map = None
+ if "AutoFeatureExtractor" in config_dict.get("auto_map", {}):
+ feature_extractor_auto_map = config_dict["auto_map"]["AutoFeatureExtractor"]
+
+ # If we don't find the feature extractor class in the feature extractor config, let's try the model config.
+ if feature_extractor_class is None and feature_extractor_auto_map is None:
+ if not isinstance(config, PreTrainedConfig):
+ config = AutoConfig.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+ # It could be in `config.feature_extractor_type``
+ feature_extractor_class = getattr(config, "feature_extractor_type", None)
+ if hasattr(config, "auto_map") and "AutoFeatureExtractor" in config.auto_map:
+ feature_extractor_auto_map = config.auto_map["AutoFeatureExtractor"]
+
+ if feature_extractor_class is not None:
+ feature_extractor_class = feature_extractor_class_from_name(feature_extractor_class)
+
+ has_remote_code = feature_extractor_auto_map is not None
+ has_local_code = feature_extractor_class is not None or type(config) in FEATURE_EXTRACTOR_MAPPING
+ explicit_local_code = has_local_code and not (
+ feature_extractor_class or FEATURE_EXTRACTOR_MAPPING[type(config)]
+ ).__module__.startswith("transformers.")
+ if has_remote_code:
+ if "--" in feature_extractor_auto_map:
+ upstream_repo = feature_extractor_auto_map.split("--")[0]
+ else:
+ upstream_repo = None
+ trust_remote_code = resolve_trust_remote_code(
+ trust_remote_code, pretrained_model_name_or_path, has_local_code, has_remote_code, upstream_repo
+ )
+
+ if has_remote_code and trust_remote_code and not explicit_local_code:
+ feature_extractor_class = get_class_from_dynamic_module(
+ feature_extractor_auto_map, pretrained_model_name_or_path, **kwargs
+ )
+ _ = kwargs.pop("code_revision", None)
+ feature_extractor_class.register_for_auto_class()
+ return feature_extractor_class.from_pretrained(pretrained_model_name_or_path, **kwargs)
+ elif feature_extractor_class is not None:
+ return feature_extractor_class.from_pretrained(pretrained_model_name_or_path, **kwargs)
+ # Last try: we use the FEATURE_EXTRACTOR_MAPPING.
+ elif type(config) in FEATURE_EXTRACTOR_MAPPING:
+ feature_extractor_class = FEATURE_EXTRACTOR_MAPPING[type(config)]
+ return feature_extractor_class.from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+ raise ValueError(
+ f"Unrecognized feature extractor in {pretrained_model_name_or_path}. Should have a "
+ f"`feature_extractor_type` key in its {FEATURE_EXTRACTOR_NAME} of {CONFIG_NAME}, or one of the following "
+ f"`model_type` keys in its {CONFIG_NAME}: {', '.join(c for c in FEATURE_EXTRACTOR_MAPPING_NAMES)}"
+ )
+
+ @staticmethod
+ def register(config_class, feature_extractor_class, exist_ok=False):
+ """
+ Register a new feature extractor for this class.
+
+ Args:
+ config_class ([`PreTrainedConfig`]):
+ The configuration corresponding to the model to register.
+ feature_extractor_class ([`FeatureExtractorMixin`]): The feature extractor to register.
+ """
+ FEATURE_EXTRACTOR_MAPPING.register(config_class, feature_extractor_class, exist_ok=exist_ok)
+
+
+__all__ = ["FEATURE_EXTRACTOR_MAPPING", "AutoFeatureExtractor"]
diff --git a/third_party/transformers/src/transformers/models/auto/image_processing_auto.py b/third_party/transformers/src/transformers/models/auto/image_processing_auto.py
new file mode 100644
index 0000000000000000000000000000000000000000..4ada3ba6b8edeab212f3814f6c8075c7adf7c956
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/auto/image_processing_auto.py
@@ -0,0 +1,819 @@
+# Copyright 2022 The HuggingFace Inc. team.
+#
+# 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.
+"""AutoImageProcessor class."""
+
+import importlib
+import os
+from collections import OrderedDict
+from typing import TYPE_CHECKING
+
+# Build the list of all image processors
+from ...configuration_utils import PreTrainedConfig
+from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
+from ...image_processing_utils import ImageProcessingMixin
+from ...utils import (
+ CONFIG_NAME,
+ IMAGE_PROCESSOR_NAME,
+ PROCESSOR_NAME,
+ cached_file,
+ is_timm_config_dict,
+ is_timm_local_checkpoint,
+ is_torchvision_available,
+ logging,
+ safe_load_json_file,
+)
+from ...utils.import_utils import requires
+from .auto_factory import _LazyAutoMapping
+from .configuration_auto import (
+ CONFIG_MAPPING_NAMES,
+ AutoConfig,
+ model_type_to_module_name,
+ replace_list_option_in_docstrings,
+)
+
+
+logger = logging.get_logger(__name__)
+
+# These image processors use Lanczos interpolation, which is not supported by fast image processors.
+# To avoid important differences in outputs, we default to using the PIL backend for these processors.
+DEFAULT_TO_PIL_BACKEND_IMAGE_PROCESSORS = [
+ "ChameleonImageProcessor",
+ "FlavaImageProcessor",
+ "Idefics3ImageProcessor",
+ "SmolVLMImageProcessor",
+]
+
+if TYPE_CHECKING:
+ # This significantly improves completion suggestion performance when
+ # the transformers package is used with Microsoft's Pylance language server.
+ IMAGE_PROCESSOR_MAPPING_NAMES: OrderedDict[str, dict[str, str | None]] = OrderedDict()
+else:
+ IMAGE_PROCESSOR_MAPPING_NAMES = OrderedDict(
+ [
+ ("aimv2", {"torchvision": "CLIPImageProcessor", "pil": "CLIPImageProcessorPil"}),
+ ("aimv2_vision_model", {"torchvision": "CLIPImageProcessor", "pil": "CLIPImageProcessorPil"}),
+ ("align", {"torchvision": "EfficientNetImageProcessor", "pil": "EfficientNetImageProcessorPil"}),
+ ("altclip", {"torchvision": "CLIPImageProcessor", "pil": "CLIPImageProcessorPil"}),
+ ("aria", {"torchvision": "AriaImageProcessor", "pil": "AriaImageProcessorPil"}),
+ ("aya_vision", {"torchvision": "GotOcr2ImageProcessor", "pil": "GotOcr2ImageProcessorPil"}),
+ ("beit", {"torchvision": "BeitImageProcessor", "pil": "BeitImageProcessorPil"}),
+ ("bit", {"torchvision": "BitImageProcessor", "pil": "BitImageProcessorPil"}),
+ ("blip", {"torchvision": "BlipImageProcessor", "pil": "BlipImageProcessorPil"}),
+ ("blip-2", {"torchvision": "BlipImageProcessor", "pil": "BlipImageProcessorPil"}),
+ ("bridgetower", {"torchvision": "BridgeTowerImageProcessor", "pil": "BridgeTowerImageProcessorPil"}),
+ ("chameleon", {"torchvision": "ChameleonImageProcessor", "pil": "ChameleonImageProcessorPil"}),
+ ("chinese_clip", {"torchvision": "ChineseCLIPImageProcessor", "pil": "ChineseCLIPImageProcessorPil"}),
+ ("chmv2", {"torchvision": "CHMv2ImageProcessor"}),
+ ("clip", {"torchvision": "CLIPImageProcessor", "pil": "CLIPImageProcessorPil"}),
+ ("clipseg", {"torchvision": "ViTImageProcessor", "pil": "ViTImageProcessorPil"}),
+ ("cohere2_vision", {"torchvision": "Cohere2VisionImageProcessor"}),
+ ("colpali", {"torchvision": "SiglipImageProcessor", "pil": "SiglipImageProcessorPil"}),
+ ("colqwen2", {"torchvision": "Qwen2VLImageProcessor", "pil": "Qwen2VLImageProcessorPil"}),
+ (
+ "conditional_detr",
+ {"torchvision": "ConditionalDetrImageProcessor", "pil": "ConditionalDetrImageProcessorPil"},
+ ),
+ ("convnext", {"torchvision": "ConvNextImageProcessor", "pil": "ConvNextImageProcessorPil"}),
+ ("convnextv2", {"torchvision": "ConvNextImageProcessor", "pil": "ConvNextImageProcessorPil"}),
+ ("cvt", {"torchvision": "ConvNextImageProcessor", "pil": "ConvNextImageProcessorPil"}),
+ ("data2vec-vision", {"torchvision": "BeitImageProcessor", "pil": "BeitImageProcessorPil"}),
+ ("deepseek_vl", {"torchvision": "DeepseekVLImageProcessor", "pil": "DeepseekVLImageProcessorPil"}),
+ (
+ "deepseek_vl_hybrid",
+ {"torchvision": "DeepseekVLHybridImageProcessor", "pil": "DeepseekVLHybridImageProcessorPil"},
+ ),
+ (
+ "deformable_detr",
+ {"torchvision": "DeformableDetrImageProcessor", "pil": "DeformableDetrImageProcessorPil"},
+ ),
+ ("deit", {"torchvision": "DeiTImageProcessor", "pil": "DeiTImageProcessorPil"}),
+ ("depth_anything", {"torchvision": "DPTImageProcessor", "pil": "DPTImageProcessorPil"}),
+ ("depth_pro", {"torchvision": "DepthProImageProcessor"}),
+ ("detr", {"torchvision": "DetrImageProcessor", "pil": "DetrImageProcessorPil"}),
+ ("dinat", {"torchvision": "ViTImageProcessor", "pil": "ViTImageProcessorPil"}),
+ ("dinov2", {"torchvision": "BitImageProcessor", "pil": "BitImageProcessorPil"}),
+ ("dinov3_vit", {"torchvision": "DINOv3ViTImageProcessor"}),
+ ("donut-swin", {"torchvision": "DonutImageProcessor", "pil": "DonutImageProcessorPil"}),
+ ("dpt", {"torchvision": "DPTImageProcessor", "pil": "DPTImageProcessorPil"}),
+ ("edgetam", {"torchvision": "Sam2ImageProcessor"}),
+ (
+ "efficientloftr",
+ {"torchvision": "EfficientLoFTRImageProcessor", "pil": "EfficientLoFTRImageProcessorPil"},
+ ),
+ ("efficientnet", {"torchvision": "EfficientNetImageProcessor", "pil": "EfficientNetImageProcessorPil"}),
+ ("emu3", {"pil": "Emu3ImageProcessor"}),
+ ("eomt", {"torchvision": "EomtImageProcessor", "pil": "EomtImageProcessorPil"}),
+ ("eomt_dinov3", {"torchvision": "EomtImageProcessor", "pil": "EomtImageProcessorPil"}),
+ (
+ "ernie4_5_vl_moe",
+ {"torchvision": "Ernie4_5_VLMoeImageProcessor", "pil": "Ernie4_5_VLMoeImageProcessorPil"},
+ ),
+ ("flava", {"torchvision": "FlavaImageProcessor", "pil": "FlavaImageProcessorPil"}),
+ ("florence2", {"torchvision": "CLIPImageProcessor", "pil": "CLIPImageProcessorPil"}),
+ ("focalnet", {"torchvision": "BitImageProcessor", "pil": "BitImageProcessorPil"}),
+ ("fuyu", {"torchvision": "FuyuImageProcessor", "pil": "FuyuImageProcessorPil"}),
+ ("gemma3", {"torchvision": "Gemma3ImageProcessor", "pil": "Gemma3ImageProcessorPil"}),
+ ("gemma3n", {"torchvision": "SiglipImageProcessor", "pil": "SiglipImageProcessorPil"}),
+ ("gemma4", {"torchvision": "Gemma4ImageProcessor", "pil": "Gemma4ImageProcessorPil"}),
+ ("git", {"torchvision": "CLIPImageProcessor", "pil": "CLIPImageProcessorPil"}),
+ ("glm46v", {"torchvision": "Glm46VImageProcessor", "pil": "Glm46VImageProcessorPil"}),
+ ("glm4v", {"torchvision": "Glm4vImageProcessor", "pil": "Glm4vImageProcessorPil"}),
+ ("glm_image", {"torchvision": "GlmImageImageProcessor", "pil": "GlmImageImageProcessorPil"}),
+ ("glpn", {"torchvision": "GLPNImageProcessor", "pil": "GLPNImageProcessorPil"}),
+ ("got_ocr2", {"torchvision": "GotOcr2ImageProcessor", "pil": "GotOcr2ImageProcessorPil"}),
+ (
+ "grounding-dino",
+ {"torchvision": "GroundingDinoImageProcessor", "pil": "GroundingDinoImageProcessorPil"},
+ ),
+ ("groupvit", {"torchvision": "CLIPImageProcessor", "pil": "CLIPImageProcessorPil"}),
+ ("hiera", {"torchvision": "BitImageProcessor", "pil": "BitImageProcessorPil"}),
+ ("idefics", {"torchvision": "IdeficsImageProcessor", "pil": "IdeficsImageProcessorPil"}),
+ ("idefics2", {"torchvision": "Idefics2ImageProcessor", "pil": "Idefics2ImageProcessorPil"}),
+ ("idefics3", {"torchvision": "Idefics3ImageProcessor", "pil": "Idefics3ImageProcessorPil"}),
+ ("ijepa", {"torchvision": "ViTImageProcessor", "pil": "ViTImageProcessorPil"}),
+ ("imagegpt", {"torchvision": "ImageGPTImageProcessor", "pil": "ImageGPTImageProcessorPil"}),
+ ("instructblip", {"torchvision": "BlipImageProcessor", "pil": "BlipImageProcessorPil"}),
+ ("internvl", {"torchvision": "GotOcr2ImageProcessor", "pil": "GotOcr2ImageProcessorPil"}),
+ ("janus", {"torchvision": "JanusImageProcessor", "pil": "JanusImageProcessorPil"}),
+ ("kosmos-2", {"torchvision": "CLIPImageProcessor", "pil": "CLIPImageProcessorPil"}),
+ ("kosmos-2.5", {"torchvision": "Kosmos2_5ImageProcessor", "pil": "Kosmos2_5ImageProcessorPil"}),
+ ("layoutlmv2", {"torchvision": "LayoutLMv2ImageProcessor", "pil": "LayoutLMv2ImageProcessorPil"}),
+ ("layoutlmv3", {"torchvision": "LayoutLMv3ImageProcessor", "pil": "LayoutLMv3ImageProcessorPil"}),
+ ("layoutxlm", {"torchvision": "LayoutLMv2ImageProcessor", "pil": "LayoutLMv2ImageProcessorPil"}),
+ ("levit", {"torchvision": "LevitImageProcessor", "pil": "LevitImageProcessorPil"}),
+ ("lfm2_vl", {"torchvision": "Lfm2VlImageProcessor"}),
+ ("lightglue", {"torchvision": "LightGlueImageProcessor", "pil": "LightGlueImageProcessorPil"}),
+ ("lighton_ocr", {"torchvision": "PixtralImageProcessor", "pil": "PixtralImageProcessorPil"}),
+ ("llama4", {"torchvision": "Llama4ImageProcessor"}),
+ ("llava", {"torchvision": "LlavaImageProcessor", "pil": "LlavaImageProcessorPil"}),
+ ("llava_next", {"torchvision": "LlavaNextImageProcessor", "pil": "LlavaNextImageProcessorPil"}),
+ ("llava_next_video", {"torchvision": "LlavaNextImageProcessor", "pil": "LlavaNextImageProcessorPil"}),
+ (
+ "llava_onevision",
+ {"torchvision": "LlavaOnevisionImageProcessor", "pil": "LlavaOnevisionImageProcessorPil"},
+ ),
+ ("lw_detr", {"torchvision": "DeformableDetrImageProcessor", "pil": "DeformableDetrImageProcessorPil"}),
+ ("mask2former", {"torchvision": "Mask2FormerImageProcessor", "pil": "Mask2FormerImageProcessorPil"}),
+ ("maskformer", {"torchvision": "MaskFormerImageProcessor", "pil": "MaskFormerImageProcessorPil"}),
+ ("metaclip_2", {"torchvision": "CLIPImageProcessor", "pil": "CLIPImageProcessorPil"}),
+ ("mgp-str", {"torchvision": "ViTImageProcessor", "pil": "ViTImageProcessorPil"}),
+ ("mistral3", {"torchvision": "PixtralImageProcessor", "pil": "PixtralImageProcessorPil"}),
+ ("mlcd", {"torchvision": "CLIPImageProcessor", "pil": "CLIPImageProcessorPil"}),
+ ("mllama", {"torchvision": "MllamaImageProcessor", "pil": "MllamaImageProcessorPil"}),
+ (
+ "mm-grounding-dino",
+ {"torchvision": "GroundingDinoImageProcessor", "pil": "GroundingDinoImageProcessorPil"},
+ ),
+ ("mobilenet_v1", {"torchvision": "MobileNetV1ImageProcessor", "pil": "MobileNetV1ImageProcessorPil"}),
+ ("mobilenet_v2", {"torchvision": "MobileNetV2ImageProcessor", "pil": "MobileNetV2ImageProcessorPil"}),
+ ("mobilevit", {"torchvision": "MobileViTImageProcessor", "pil": "MobileViTImageProcessorPil"}),
+ ("mobilevitv2", {"torchvision": "MobileViTImageProcessor", "pil": "MobileViTImageProcessorPil"}),
+ ("nougat", {"torchvision": "NougatImageProcessor", "pil": "NougatImageProcessorPil"}),
+ ("omdet-turbo", {"torchvision": "DetrImageProcessor", "pil": "DetrImageProcessorPil"}),
+ ("oneformer", {"torchvision": "OneFormerImageProcessor", "pil": "OneFormerImageProcessorPil"}),
+ ("ovis2", {"torchvision": "Ovis2ImageProcessor", "pil": "Ovis2ImageProcessorPil"}),
+ ("owlv2", {"torchvision": "Owlv2ImageProcessor", "pil": "Owlv2ImageProcessorPil"}),
+ ("owlvit", {"torchvision": "OwlViTImageProcessor", "pil": "OwlViTImageProcessorPil"}),
+ ("paddleocr_vl", {"torchvision": "PaddleOCRVLImageProcessor", "pil": "PaddleOCRVLImageProcessorPil"}),
+ ("paligemma", {"torchvision": "SiglipImageProcessor", "pil": "SiglipImageProcessorPil"}),
+ ("perceiver", {"torchvision": "PerceiverImageProcessor", "pil": "PerceiverImageProcessorPil"}),
+ ("perception_lm", {"torchvision": "PerceptionLMImageProcessor"}),
+ ("phi4_multimodal", {"torchvision": "Phi4MultimodalImageProcessor"}),
+ ("pi0", {"torchvision": "PI0ImageProcessor"}),
+ ("pix2struct", {"torchvision": "Pix2StructImageProcessor", "pil": "Pix2StructImageProcessorPil"}),
+ ("pixio", {"torchvision": "BitImageProcessor", "pil": "BitImageProcessorPil"}),
+ ("pixtral", {"torchvision": "PixtralImageProcessor", "pil": "PixtralImageProcessorPil"}),
+ ("poolformer", {"torchvision": "PoolFormerImageProcessor", "pil": "PoolFormerImageProcessorPil"}),
+ (
+ "pp_chart2table",
+ {"torchvision": "PPChart2TableImageProcessor", "pil": "PPChart2TableImageProcessorPil"},
+ ),
+ ("pp_doclayout_v2", {"torchvision": "PPDocLayoutV2ImageProcessor"}),
+ ("pp_doclayout_v3", {"torchvision": "PPDocLayoutV3ImageProcessor"}),
+ ("pp_lcnet", {"torchvision": "PPLCNetImageProcessor"}),
+ ("pp_ocrv5_mobile_det", {"torchvision": "PPOCRV5ServerDetImageProcessor"}),
+ ("pp_ocrv5_mobile_rec", {"torchvision": "PPOCRV5ServerRecImageProcessor"}),
+ ("pp_ocrv5_server_det", {"torchvision": "PPOCRV5ServerDetImageProcessor"}),
+ ("pp_ocrv5_server_rec", {"torchvision": "PPOCRV5ServerRecImageProcessor"}),
+ (
+ "prompt_depth_anything",
+ {"torchvision": "PromptDepthAnythingImageProcessor", "pil": "PromptDepthAnythingImageProcessorPil"},
+ ),
+ ("pvt", {"torchvision": "PvtImageProcessor", "pil": "PvtImageProcessorPil"}),
+ ("pvt_v2", {"torchvision": "PvtImageProcessor", "pil": "PvtImageProcessorPil"}),
+ ("qwen2_5_omni", {"torchvision": "Qwen2VLImageProcessor", "pil": "Qwen2VLImageProcessorPil"}),
+ ("qwen2_5_vl", {"torchvision": "Qwen2VLImageProcessor", "pil": "Qwen2VLImageProcessorPil"}),
+ ("qwen2_vl", {"torchvision": "Qwen2VLImageProcessor", "pil": "Qwen2VLImageProcessorPil"}),
+ ("qwen3_5", {"torchvision": "Qwen2VLImageProcessor", "pil": "Qwen2VLImageProcessorPil"}),
+ ("qwen3_5_moe", {"torchvision": "Qwen2VLImageProcessor", "pil": "Qwen2VLImageProcessorPil"}),
+ ("qwen3_omni_moe", {"torchvision": "Qwen2VLImageProcessor", "pil": "Qwen2VLImageProcessorPil"}),
+ ("qwen3_vl", {"torchvision": "Qwen2VLImageProcessor", "pil": "Qwen2VLImageProcessorPil"}),
+ ("regnet", {"torchvision": "ConvNextImageProcessor", "pil": "ConvNextImageProcessorPil"}),
+ ("resnet", {"torchvision": "ConvNextImageProcessor", "pil": "ConvNextImageProcessorPil"}),
+ ("rt_detr", {"torchvision": "RTDetrImageProcessor", "pil": "RTDetrImageProcessorPil"}),
+ ("sam", {"torchvision": "SamImageProcessor", "pil": "SamImageProcessorPil"}),
+ ("sam2", {"torchvision": "Sam2ImageProcessor"}),
+ ("sam2_video", {"torchvision": "Sam2ImageProcessor"}),
+ ("sam3", {"torchvision": "Sam3ImageProcessor"}),
+ ("sam3_tracker", {"torchvision": "Sam3ImageProcessor"}),
+ ("sam3_tracker_video", {"torchvision": "Sam3ImageProcessor"}),
+ ("sam3_video", {"torchvision": "Sam3ImageProcessor"}),
+ ("sam_hq", {"torchvision": "SamImageProcessor", "pil": "SamImageProcessorPil"}),
+ ("segformer", {"torchvision": "SegformerImageProcessor", "pil": "SegformerImageProcessorPil"}),
+ ("seggpt", {"torchvision": "SegGptImageProcessor", "pil": "SegGptImageProcessorPil"}),
+ ("shieldgemma2", {"torchvision": "Gemma3ImageProcessor", "pil": "Gemma3ImageProcessorPil"}),
+ ("siglip", {"torchvision": "SiglipImageProcessor", "pil": "SiglipImageProcessorPil"}),
+ ("siglip2", {"torchvision": "Siglip2ImageProcessor", "pil": "Siglip2ImageProcessorPil"}),
+ ("slanext", {"torchvision": "SLANeXtImageProcessor"}),
+ ("smolvlm", {"torchvision": "SmolVLMImageProcessor", "pil": "SmolVLMImageProcessorPil"}),
+ ("superglue", {"torchvision": "SuperGlueImageProcessor", "pil": "SuperGlueImageProcessorPil"}),
+ ("superpoint", {"torchvision": "SuperPointImageProcessor", "pil": "SuperPointImageProcessorPil"}),
+ ("swiftformer", {"torchvision": "ViTImageProcessor", "pil": "ViTImageProcessorPil"}),
+ ("swin", {"torchvision": "ViTImageProcessor", "pil": "ViTImageProcessorPil"}),
+ ("swin2sr", {"torchvision": "Swin2SRImageProcessor", "pil": "Swin2SRImageProcessorPil"}),
+ ("swinv2", {"torchvision": "ViTImageProcessor", "pil": "ViTImageProcessorPil"}),
+ ("t5gemma2", {"torchvision": "Gemma3ImageProcessor", "pil": "Gemma3ImageProcessorPil"}),
+ ("t5gemma2_encoder", {"torchvision": "Gemma3ImageProcessor", "pil": "Gemma3ImageProcessorPil"}),
+ ("table-transformer", {"torchvision": "DetrImageProcessor", "pil": "DetrImageProcessorPil"}),
+ ("textnet", {"torchvision": "TextNetImageProcessor", "pil": "TextNetImageProcessorPil"}),
+ ("timesformer", {"pil": "VideoMAEImageProcessorPil", "torchvision": "VideoMAEImageProcessor"}),
+ ("timm_wrapper", {"pil": "TimmWrapperImageProcessor"}),
+ ("trocr", {"torchvision": "ViTImageProcessor", "pil": "ViTImageProcessorPil"}),
+ ("tvp", {"torchvision": "TvpImageProcessor", "pil": "TvpImageProcessorPil"}),
+ ("udop", {"torchvision": "LayoutLMv3ImageProcessor", "pil": "LayoutLMv3ImageProcessorPil"}),
+ ("upernet", {"torchvision": "SegformerImageProcessor", "pil": "SegformerImageProcessorPil"}),
+ ("uvdoc", {"torchvision": "UVDocImageProcessor"}),
+ ("video_llama_3", {"torchvision": "VideoLlama3ImageProcessor", "pil": "VideoLlama3ImageProcessorPil"}),
+ ("video_llava", {"pil": "VideoLlavaImageProcessor"}),
+ ("videomae", {"torchvision": "VideoMAEImageProcessor", "pil": "VideoMAEImageProcessorPil"}),
+ ("vilt", {"torchvision": "ViltImageProcessor", "pil": "ViltImageProcessorPil"}),
+ ("vipllava", {"torchvision": "CLIPImageProcessor", "pil": "CLIPImageProcessorPil"}),
+ ("vit", {"torchvision": "ViTImageProcessor", "pil": "ViTImageProcessorPil"}),
+ ("vit_mae", {"torchvision": "ViTImageProcessor", "pil": "ViTImageProcessorPil"}),
+ ("vit_msn", {"torchvision": "ViTImageProcessor", "pil": "ViTImageProcessorPil"}),
+ ("vitmatte", {"torchvision": "VitMatteImageProcessor", "pil": "VitMatteImageProcessorPil"}),
+ ("vitpose", {"torchvision": "VitPoseImageProcessor", "pil": "VitPoseImageProcessorPil"}),
+ ("vivit", {"torchvision": "VivitImageProcessor"}),
+ ("xclip", {"torchvision": "CLIPImageProcessor", "pil": "CLIPImageProcessorPil"}),
+ ("yolos", {"torchvision": "YolosImageProcessor", "pil": "YolosImageProcessorPil"}),
+ ("zoedepth", {"torchvision": "ZoeDepthImageProcessor", "pil": "ZoeDepthImageProcessorPil"}),
+ ]
+ )
+
+IMAGE_PROCESSOR_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, IMAGE_PROCESSOR_MAPPING_NAMES)
+
+
+def get_image_processor_class_from_name(class_name: str):
+ """Resolve an image processor class name to its class. Handles both base names (e.g. CLIPImageProcessor)
+ and PIL backend names (e.g. CLIPImageProcessorPil). No recursion needed since names are direct."""
+ if class_name == "BaseImageProcessorFast":
+ # kept for backward compatibility - return TorchvisionBackend
+ from ...image_processing_backends import TorchvisionBackend
+
+ return TorchvisionBackend
+
+ # First, check registered extra content (user-registered classes)
+ for mapping in IMAGE_PROCESSOR_MAPPING._extra_content.values():
+ for extractor_class in mapping.values():
+ if isinstance(extractor_class, type) and getattr(extractor_class, "__name__", None) == class_name:
+ return extractor_class
+
+ # Check the mapping names - class names are either base (torchvision) or base+Pil (pil)
+ for model_type, extractors_dict in IMAGE_PROCESSOR_MAPPING_NAMES.items():
+ if class_name in extractors_dict.values():
+ module_name = model_type_to_module_name(model_type)
+ module = importlib.import_module(f".{module_name}", "transformers.models")
+ try:
+ return getattr(module, class_name)
+ except AttributeError:
+ continue
+
+ # Fallback: class may be in main init (e.g. when dep is missing, returns dummy)
+ main_module = importlib.import_module("transformers")
+ if hasattr(main_module, class_name):
+ return getattr(main_module, class_name)
+
+ return None
+
+
+def get_image_processor_config(
+ pretrained_model_name_or_path: str | os.PathLike,
+ cache_dir: str | os.PathLike | None = None,
+ force_download: bool = False,
+ proxies: dict[str, str] | None = None,
+ token: bool | str | None = None,
+ revision: str | None = None,
+ local_files_only: bool = False,
+ **kwargs,
+):
+ """
+ Loads the image processor configuration from a pretrained model image processor configuration.
+
+ Args:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained model configuration hosted inside a model repo on
+ huggingface.co.
+ - a path to a *directory* containing a configuration file saved using the
+ [`~ProcessorMixin.save_pretrained`] method, e.g., `./my_model_directory/`.
+
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
+ cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the configuration files and override the cached versions if they
+ exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `hf auth login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ local_files_only (`bool`, *optional*, defaults to `False`):
+ If `True`, will only try to load the image processor configuration from local files.
+
+
+
+ Passing `token=True` is required when you want to use a private model.
+
+
+
+ Returns:
+ `Dict`: The configuration of the image processor.
+
+ Examples:
+
+ ```python
+ # Download configuration from huggingface.co and cache.
+ image_processor_config = get_image_processor_config("google-bert/bert-base-uncased")
+ # This model does not have a image processor config so the result will be an empty dict.
+ image_processor_config = get_image_processor_config("FacebookAI/xlm-roberta-base")
+
+ # Save a pretrained image processor locally and you can reload its config
+ from transformers import AutoImageProcessor
+
+ image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
+ image_processor.save_pretrained("image-processor-test")
+ image_processor_config = get_image_processor_config("image-processor-test")
+ ```"""
+ # Load with a priority given to the nested processor config, if available in repo
+ resolved_processor_file = cached_file(
+ pretrained_model_name_or_path,
+ filename=PROCESSOR_NAME,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ token=token,
+ revision=revision,
+ local_files_only=local_files_only,
+ _raise_exceptions_for_gated_repo=False,
+ _raise_exceptions_for_missing_entries=False,
+ )
+ resolved_image_processor_file = cached_file(
+ pretrained_model_name_or_path,
+ filename=IMAGE_PROCESSOR_NAME,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ token=token,
+ revision=revision,
+ local_files_only=local_files_only,
+ _raise_exceptions_for_gated_repo=False,
+ _raise_exceptions_for_missing_entries=False,
+ )
+
+ # An empty list if none of the possible files is found in the repo
+ if not resolved_image_processor_file and not resolved_processor_file:
+ logger.info("Could not locate the image processor configuration file.")
+ return {}
+
+ # Load image_processor dict. Priority goes as (nested config if found -> image processor config)
+ # We are downloading both configs because almost all models have a `processor_config.json` but
+ # not all of these are nested. We need to check if it was saved recently as nested or if it is legacy style
+ image_processor_dict = {}
+ if resolved_processor_file is not None:
+ processor_dict = safe_load_json_file(resolved_processor_file)
+ if "image_processor" in processor_dict:
+ image_processor_dict = processor_dict["image_processor"]
+
+ if resolved_image_processor_file is not None and image_processor_dict is None:
+ image_processor_dict = safe_load_json_file(resolved_image_processor_file)
+
+ return image_processor_dict
+
+
+def _resolve_backend(backend: str | None, use_fast: bool | None, base_class_name: str | None) -> str:
+ """Resolve raw backend inputs to a concrete backend name ('torchvision' or 'pil').
+
+ Handles, in order:
+ - Deprecated ``use_fast`` flag: warns and converts to an explicit backend string when no
+ explicit backend is given.
+ - Explicit backend string: returned as-is.
+ - None resolution: forces 'pil' for processors in DEFAULT_TO_PIL_BACKEND_IMAGE_PROCESSORS
+ (Lanczos interpolation, unsupported by torchvision); otherwise picks 'torchvision' when
+ available, falling back to 'pil'.
+ """
+ if use_fast is not None:
+ logger.warning_once(
+ "The `use_fast` parameter is deprecated and will be removed in a future version. "
+ 'Use `backend="torchvision"` instead of `use_fast=True`, or `backend="pil"` instead of `use_fast=False`.'
+ )
+ if backend is None:
+ backend = "torchvision" if use_fast else "pil"
+
+ if backend is None:
+ if base_class_name in DEFAULT_TO_PIL_BACKEND_IMAGE_PROCESSORS:
+ return "pil"
+ return "torchvision" if is_torchvision_available() else "pil"
+
+ return backend
+
+
+def _load_class_with_fallback(mapping, backend):
+ """
+ Load an image processor class from a backend-to-class mapping, with fallback.
+
+ Tries the requested backend first, then the opposite standard backend,
+ then any remaining backends. Works with both string class names and resolved class objects.
+
+ Unavailable backends are detected via DummyObject: classes whose required libraries are missing
+ are represented as DummyObject subclasses (is_dummy=True). When the torchvision backend is
+ missing but a PIL variant exists, _LazyModule transparently returns the PIL class with its own
+ warning, so _load_class_with_fallback naturally receives a usable class without extra gating.
+
+ Args:
+ mapping: dict mapping backend names (str) to class names (str) or class objects (type).
+ backend: the preferred backend name (e.g. "torchvision", "pil").
+
+ Returns:
+ The loaded class, or None if no class could be loaded.
+ """
+ backends_to_try = [backend] + [k for k in mapping if k != backend]
+
+ for b in backends_to_try:
+ value = mapping.get(b)
+ if value is None:
+ continue
+
+ # Value can be a class object (from resolved mapping) or a string class name
+ if isinstance(value, type):
+ processor_class = value
+ else:
+ processor_class = get_image_processor_class_from_name(value)
+
+ if processor_class is None or getattr(processor_class, "is_dummy", False):
+ continue
+
+ if b != backend:
+ logger.warning_once(f"Requested {backend} backend is not available. Falling back to {b} backend.")
+ return processor_class
+
+ return None
+
+
+def _find_mapping_for_image_processor(base_class_name: str) -> dict | None:
+ """
+ Find the backend->class mapping that contains base_class_name in its values.
+ Returns the mapping dict (including any custom registered backends) or None.
+ """
+
+ def _value_matches(val, name: str) -> bool:
+ if val is None:
+ return False
+ if isinstance(val, str):
+ return val == name
+ if isinstance(val, type):
+ return getattr(val, "__name__", None) == name
+ return False
+
+ for mapping_dict in IMAGE_PROCESSOR_MAPPING_NAMES.values():
+ if any(_value_matches(v, base_class_name) for v in mapping_dict.values()):
+ return mapping_dict
+
+ for content in IMAGE_PROCESSOR_MAPPING._extra_content.values():
+ if any(_value_matches(v, base_class_name) for v in content.values()):
+ return content
+
+ return None
+
+
+def _load_backend_class(base_class_name, backend, is_legacy_fast=False):
+ """
+ Load image processor class for a given backend. Uses the mapping from
+ IMAGE_PROCESSOR_MAPPING when base_class_name is found in its values (so config
+ overrides and custom backends are respected). Falls back to base+Pil convention
+ for remote code / unknown processors.
+ """
+ mapping = _find_mapping_for_image_processor(base_class_name)
+ if mapping is None:
+ mapping = {
+ "torchvision": base_class_name,
+ "pil": base_class_name + "Pil",
+ }
+ processor_class = _load_class_with_fallback(mapping, backend)
+
+ # For legacy Fast classes, try the original Fast class name as last resort
+ if processor_class is None and is_legacy_fast:
+ processor_class = get_image_processor_class_from_name(base_class_name + "Fast")
+
+ return processor_class
+
+
+def _resolve_auto_map_class_ref(auto_map, backend):
+ """Extract the class reference string from an auto_map entry based on backend preference.
+
+ Returns:
+ A string that may be:
+ - A simple class name (e.g. `"MyImageProcessor"`)
+ - A Hub reference in the form `upstream_repo--path/to/file.py::ClassName`, where the part before
+ `--` is the upstream repo ID (used for trust_remote_code resolution).
+ """
+ if isinstance(auto_map, dict):
+ return auto_map.get(backend) or next(iter(auto_map.values()))
+ if isinstance(auto_map, (list, tuple)):
+ if backend == "torchvision" and len(auto_map) > 1 and auto_map[1] is not None:
+ return auto_map[1]
+ return auto_map[0]
+ # Single string (legacy)
+ return auto_map
+
+
+@requires(backends=("vision",))
+class AutoImageProcessor:
+ r"""
+ This is a generic image processor class that will be instantiated as one of the image processor classes of the
+ library when created with the [`AutoImageProcessor.from_pretrained`] class method.
+
+ This class cannot be instantiated directly using `__init__()` (throws an error).
+ """
+
+ def __init__(self):
+ raise OSError(
+ "AutoImageProcessor is designed to be instantiated "
+ "using the `AutoImageProcessor.from_pretrained(pretrained_model_name_or_path)` method."
+ )
+
+ @classmethod
+ @replace_list_option_in_docstrings(IMAGE_PROCESSOR_MAPPING_NAMES)
+ def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
+ r"""
+ Instantiate one of the image processor classes of the library from a pretrained model vocabulary.
+
+ The image processor class to instantiate is selected based on the `model_type` property of the config object
+ (either passed as an argument or loaded from `pretrained_model_name_or_path` if possible), or when it's
+ missing, by falling back to using pattern matching on `pretrained_model_name_or_path`:
+
+ List options
+
+ Params:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained image_processor hosted inside a model repo on
+ huggingface.co.
+ - a path to a *directory* containing a image processor file saved using the
+ [`~image_processing_utils.ImageProcessingMixin.save_pretrained`] method, e.g.,
+ `./my_model_directory/`.
+ - a path to a saved image processor JSON *file*, e.g.,
+ `./my_model_directory/preprocessor_config.json`.
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model image processor should be cached if the
+ standard cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the image processor files and override the cached versions if
+ they exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `hf auth login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ use_fast (`bool`, *optional*, defaults to `False`):
+ **Deprecated**: Use `backend="torchvision"` instead. This parameter is kept for backward compatibility.
+ Use a fast torchvision-based image processor if it is supported for a given model.
+ If a fast image processor is not available for a given model, a normal numpy-based image processor
+ is returned instead.
+ backend (`str`, *optional*, defaults to `None`):
+ The backend to use for image processing. Can be:
+ - `None`: Automatically select the best available backend (torchvision if available, otherwise pil)
+ - `"torchvision"`: Use Torchvision backend (GPU-accelerated, faster)
+ - `"pil"`: Use PIL backend (portable, CPU-only)
+ - Any custom backend name registered via `register()` method
+ return_unused_kwargs (`bool`, *optional*, defaults to `False`):
+ If `False`, then this function returns just the final image processor object. If `True`, then this
+ functions returns a `Tuple(image_processor, unused_kwargs)` where *unused_kwargs* is a dictionary
+ consisting of the key/value pairs whose keys are not image processor attributes: i.e., the part of
+ `kwargs` which has not been used to update `image_processor` and is otherwise ignored.
+ trust_remote_code (`bool`, *optional*, defaults to `False`):
+ Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
+ should only be set to `True` for repositories you trust and in which you have read the code, as it will
+ execute code present on the Hub on your local machine.
+ image_processor_filename (`str`, *optional*, defaults to `"config.json"`):
+ The name of the file in the model directory to use for the image processor config.
+ kwargs (`dict[str, Any]`, *optional*):
+ The values in kwargs of any keys which are image processor attributes will be used to override the
+ loaded values. Behavior concerning key/value pairs whose keys are *not* image processor attributes is
+ controlled by the `return_unused_kwargs` keyword parameter.
+
+
+
+ Passing `token=True` is required when you want to use a private model.
+
+
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoImageProcessor
+
+ >>> # Download image processor from huggingface.co and cache.
+ >>> image_processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
+
+ >>> # If image processor files are in a directory (e.g. image processor was saved using *save_pretrained('./test/saved_model/')*)
+ >>> # image_processor = AutoImageProcessor.from_pretrained("./test/saved_model/")
+ ```"""
+ config = kwargs.pop("config", None)
+ use_fast = kwargs.pop("use_fast", None)
+ backend_kwarg = kwargs.pop("backend", None)
+ trust_remote_code = kwargs.pop("trust_remote_code", None)
+ kwargs["_from_auto"] = True
+
+ # Resolve the image processor config filename
+ if "image_processor_filename" in kwargs:
+ image_processor_filename = kwargs.pop("image_processor_filename")
+ elif is_timm_local_checkpoint(pretrained_model_name_or_path):
+ image_processor_filename = CONFIG_NAME
+ else:
+ image_processor_filename = IMAGE_PROCESSOR_NAME
+
+ # Load the image processor config
+
+ try:
+ config_dict, _ = ImageProcessingMixin.get_image_processor_dict(
+ pretrained_model_name_or_path, image_processor_filename=image_processor_filename, **kwargs
+ )
+ except Exception as initial_exception:
+ # Fallback for Hub TimmWrapper checkpoints (image processing in config.json, not preprocessor_config.json)
+ try:
+ config_dict, _ = ImageProcessingMixin.get_image_processor_dict(
+ pretrained_model_name_or_path, image_processor_filename=CONFIG_NAME, **kwargs
+ )
+ except Exception:
+ raise initial_exception
+
+ if not is_timm_config_dict(config_dict):
+ raise initial_exception
+
+ image_processor_type = config_dict.get("image_processor_type", None)
+ image_processor_auto_map = None
+ if "AutoImageProcessor" in config_dict.get("auto_map", {}):
+ image_processor_auto_map = config_dict["auto_map"]["AutoImageProcessor"]
+
+ # Backward compat: infer from feature extractor config
+ if image_processor_type is None and image_processor_auto_map is None:
+ feature_extractor_class = config_dict.pop("feature_extractor_type", None)
+ if feature_extractor_class is not None:
+ image_processor_type = feature_extractor_class.replace("FeatureExtractor", "ImageProcessor")
+ if "AutoFeatureExtractor" in config_dict.get("auto_map", {}):
+ feature_extractor_auto_map = config_dict["auto_map"]["AutoFeatureExtractor"]
+ image_processor_auto_map = feature_extractor_auto_map.replace("FeatureExtractor", "ImageProcessor")
+
+ # If not in image processor config, try the model config
+ if image_processor_type is None and image_processor_auto_map is None:
+ if not isinstance(config, PreTrainedConfig):
+ config = AutoConfig.from_pretrained(
+ pretrained_model_name_or_path,
+ trust_remote_code=trust_remote_code,
+ **kwargs,
+ )
+ image_processor_type = getattr(config, "image_processor_type", None)
+ if hasattr(config, "auto_map") and "AutoImageProcessor" in config.auto_map:
+ image_processor_auto_map = config.auto_map["AutoImageProcessor"]
+
+ # Derive base_class_name from image_processor_type
+ is_legacy_fast = False
+ base_class_name = None
+ if image_processor_type is not None:
+ is_legacy_fast = image_processor_type.endswith("Fast")
+ base_class_name = image_processor_type[:-4] if is_legacy_fast else image_processor_type
+
+ backend = _resolve_backend(backend_kwarg, use_fast, base_class_name)
+
+ image_processor_class = None
+ if base_class_name is not None:
+ image_processor_class = _load_backend_class(base_class_name, backend, is_legacy_fast)
+
+ # Handle remote code
+ has_remote_code = image_processor_auto_map is not None
+ has_local_code = image_processor_class is not None or type(config) in IMAGE_PROCESSOR_MAPPING
+ explicit_local_code = has_local_code and not (
+ image_processor_class or _load_class_with_fallback(IMAGE_PROCESSOR_MAPPING[type(config)], backend)
+ ).__module__.startswith("transformers.")
+ if has_remote_code:
+ class_ref = _resolve_auto_map_class_ref(image_processor_auto_map, backend)
+ upstream_repo = class_ref.split("--")[0] if "--" in class_ref else None
+ trust_remote_code = resolve_trust_remote_code(
+ trust_remote_code, pretrained_model_name_or_path, has_local_code, has_remote_code, upstream_repo
+ )
+
+ if has_remote_code and trust_remote_code and not explicit_local_code:
+ image_processor_class = get_class_from_dynamic_module(class_ref, pretrained_model_name_or_path, **kwargs)
+ _ = kwargs.pop("code_revision", None)
+ image_processor_class.register_for_auto_class()
+ return image_processor_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+ elif image_processor_class is not None:
+ return image_processor_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+ # Last try: we use the IMAGE_PROCESSOR_MAPPING.
+ elif type(config) in IMAGE_PROCESSOR_MAPPING:
+ image_processor_mapping = IMAGE_PROCESSOR_MAPPING[type(config)]
+ image_processor_class = _load_class_with_fallback(image_processor_mapping, backend)
+
+ if image_processor_class is not None:
+ return image_processor_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+
+ available = [k for k, v in image_processor_mapping.items() if v is not None]
+ raise ValueError(f"Could not find image processor class. Available backends: {', '.join(available)}")
+ raise ValueError(
+ f"Unrecognized image processor in {pretrained_model_name_or_path}. Should have a "
+ f"`image_processor_type` key in its {IMAGE_PROCESSOR_NAME} of {CONFIG_NAME}, or one of the following "
+ f"`model_type` keys in its {CONFIG_NAME}: {', '.join(c for c in IMAGE_PROCESSOR_MAPPING_NAMES)}"
+ )
+
+ @staticmethod
+ def register(
+ config_class,
+ slow_image_processor_class: type | None = None,
+ fast_image_processor_class: type | None = None,
+ image_processor_classes: dict[str, type] | None = None,
+ exist_ok: bool = False,
+ ):
+ """
+ Register a new image processor for this class.
+
+ Args:
+ config_class ([`PreTrainedConfig`]):
+ The configuration corresponding to the model to register.
+ slow_image_processor_class (`type`, *optional*):
+ The PIL backend image processor class (deprecated, use `image_processor_classes={"pil": ...}`).
+ fast_image_processor_class (`type`, *optional*):
+ The Torchvision backend image processor class (deprecated, use `image_processor_classes={"torchvision": ...}`).
+ image_processor_classes (`dict[str, type]`, *optional*):
+ Dictionary mapping backend names to image processor classes. Allows registering custom backends.
+ Example: `{"pil": MyPilProcessor, "torchvision": MyTorchvisionProcessor, "custom": MyCustomProcessor}`
+ exist_ok (`bool`, *optional*, defaults to `False`):
+ If `True`, allow overwriting existing registrations.
+ """
+ # Handle backward compatibility: convert old parameters to new format
+ if image_processor_classes is None:
+ image_processor_classes = {}
+ if slow_image_processor_class is not None:
+ image_processor_classes["pil"] = slow_image_processor_class
+ if fast_image_processor_class is not None:
+ image_processor_classes["torchvision"] = fast_image_processor_class
+
+ if not image_processor_classes:
+ raise ValueError(
+ "You need to specify at least one image processor class. "
+ "Use `image_processor_classes={'backend_name': ProcessorClass}` or the deprecated "
+ "`slow_image_processor_class`/`fast_image_processor_class` parameters."
+ )
+
+ # Avoid resetting existing processors if we are passing partial updates
+ if config_class in IMAGE_PROCESSOR_MAPPING._extra_content:
+ existing_mapping = IMAGE_PROCESSOR_MAPPING[config_class]
+ existing_mapping.update(image_processor_classes)
+ image_processor_classes = existing_mapping
+
+ # Validate that all classes are proper image processor classes
+ from ...image_processing_utils import BaseImageProcessor
+
+ for backend_key, processor_class in image_processor_classes.items():
+ if processor_class is not None and not issubclass(processor_class, BaseImageProcessor):
+ raise ValueError(
+ f"Image processor class for backend '{backend_key}' must inherit from `BaseImageProcessor`. "
+ f"Got: {processor_class}"
+ )
+ IMAGE_PROCESSOR_MAPPING.register(config_class, image_processor_classes, exist_ok=exist_ok)
+
+
+__all__ = ["IMAGE_PROCESSOR_MAPPING", "AutoImageProcessor"]
diff --git a/third_party/transformers/src/transformers/models/auto/modeling_auto.py b/third_party/transformers/src/transformers/models/auto/modeling_auto.py
new file mode 100644
index 0000000000000000000000000000000000000000..d4cb17cddfa6403a67d6834bc85a28dc8aaa855b
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/auto/modeling_auto.py
@@ -0,0 +1,2374 @@
+# Copyright 2018 The HuggingFace Inc. team.
+#
+# 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.
+"""Auto Model class."""
+
+import os
+from collections import OrderedDict
+from typing import TYPE_CHECKING
+
+from ...utils import logging
+from .auto_factory import (
+ _BaseAutoBackboneClass,
+ _BaseAutoModelClass,
+ _LazyAutoMapping,
+ auto_class_update,
+)
+from .configuration_auto import CONFIG_MAPPING_NAMES
+
+
+if TYPE_CHECKING:
+ from ...generation import GenerationMixin
+ from ...modeling_utils import PreTrainedModel
+
+ # class for better type annotations
+ class _BaseModelWithGenerate(PreTrainedModel, GenerationMixin):
+ pass
+
+
+logger = logging.get_logger(__name__)
+
+MODEL_MAPPING_NAMES = OrderedDict(
+ [
+ # Base model mapping
+ ("afmoe", "AfmoeModel"),
+ ("aimv2", "Aimv2Model"),
+ ("aimv2_vision_model", "Aimv2VisionModel"),
+ ("albert", "AlbertModel"),
+ ("align", "AlignModel"),
+ ("altclip", "AltCLIPModel"),
+ ("apertus", "ApertusModel"),
+ ("arcee", "ArceeModel"),
+ ("aria", "AriaModel"),
+ ("aria_text", "AriaTextModel"),
+ ("audio-spectrogram-transformer", "ASTModel"),
+ ("audioflamingo3", "AudioFlamingo3ForConditionalGeneration"),
+ ("audioflamingo3_encoder", "AudioFlamingo3Encoder"),
+ ("autoformer", "AutoformerModel"),
+ ("aya_vision", "AyaVisionModel"),
+ ("bamba", "BambaModel"),
+ ("bark", "BarkModel"),
+ ("bart", "BartModel"),
+ ("beit", "BeitModel"),
+ ("bert", "BertModel"),
+ ("bert-generation", "BertGenerationEncoder"),
+ ("big_bird", "BigBirdModel"),
+ ("bigbird_pegasus", "BigBirdPegasusModel"),
+ ("biogpt", "BioGptModel"),
+ ("bit", "BitModel"),
+ ("bitnet", "BitNetModel"),
+ ("blenderbot", "BlenderbotModel"),
+ ("blenderbot-small", "BlenderbotSmallModel"),
+ ("blip", "BlipModel"),
+ ("blip-2", "Blip2Model"),
+ ("blip_2_qformer", "Blip2QFormerModel"),
+ ("bloom", "BloomModel"),
+ ("blt", "BltModel"),
+ ("bridgetower", "BridgeTowerModel"),
+ ("bros", "BrosModel"),
+ ("camembert", "CamembertModel"),
+ ("canine", "CanineModel"),
+ ("chameleon", "ChameleonModel"),
+ ("chinese_clip", "ChineseCLIPModel"),
+ ("chinese_clip_vision_model", "ChineseCLIPVisionModel"),
+ ("clap", "ClapModel"),
+ ("clip", "CLIPModel"),
+ ("clip_text_model", "CLIPTextModel"),
+ ("clip_vision_model", "CLIPVisionModel"),
+ ("clipseg", "CLIPSegModel"),
+ ("clvp", "ClvpModelForConditionalGeneration"),
+ ("code_llama", "LlamaModel"),
+ ("codegen", "CodeGenModel"),
+ ("cohere", "CohereModel"),
+ ("cohere2", "Cohere2Model"),
+ ("cohere2_vision", "Cohere2VisionModel"),
+ ("cohere_asr", "CohereAsrModel"),
+ ("conditional_detr", "ConditionalDetrModel"),
+ ("convbert", "ConvBertModel"),
+ ("convnext", "ConvNextModel"),
+ ("convnextv2", "ConvNextV2Model"),
+ ("cpmant", "CpmAntModel"),
+ ("csm", "CsmForConditionalGeneration"),
+ ("ctrl", "CTRLModel"),
+ ("cvt", "CvtModel"),
+ ("cwm", "CwmModel"),
+ ("d_fine", "DFineModel"),
+ ("dab-detr", "DabDetrModel"),
+ ("dac", "DacModel"),
+ ("data2vec-audio", "Data2VecAudioModel"),
+ ("data2vec-text", "Data2VecTextModel"),
+ ("data2vec-vision", "Data2VecVisionModel"),
+ ("dbrx", "DbrxModel"),
+ ("deberta", "DebertaModel"),
+ ("deberta-v2", "DebertaV2Model"),
+ ("decision_transformer", "DecisionTransformerModel"),
+ ("deepseek_v2", "DeepseekV2Model"),
+ ("deepseek_v3", "DeepseekV3Model"),
+ ("deepseek_vl", "DeepseekVLModel"),
+ ("deepseek_vl_hybrid", "DeepseekVLHybridModel"),
+ ("deformable_detr", "DeformableDetrModel"),
+ ("deit", "DeiTModel"),
+ ("depth_pro", "DepthProModel"),
+ ("detr", "DetrModel"),
+ ("dia", "DiaModel"),
+ ("diffllama", "DiffLlamaModel"),
+ ("dinat", "DinatModel"),
+ ("dinov2", "Dinov2Model"),
+ ("dinov2_with_registers", "Dinov2WithRegistersModel"),
+ ("dinov3_convnext", "DINOv3ConvNextModel"),
+ ("dinov3_vit", "DINOv3ViTModel"),
+ ("distilbert", "DistilBertModel"),
+ ("doge", "DogeModel"),
+ ("donut-swin", "DonutSwinModel"),
+ ("dots1", "Dots1Model"),
+ ("dpr", "DPRQuestionEncoder"),
+ ("dpt", "DPTModel"),
+ ("edgetam", "EdgeTamModel"),
+ ("edgetam_video", "EdgeTamVideoModel"),
+ ("edgetam_vision_model", "EdgeTamVisionModel"),
+ ("efficientloftr", "EfficientLoFTRModel"),
+ ("efficientnet", "EfficientNetModel"),
+ ("electra", "ElectraModel"),
+ ("emu3", "Emu3Model"),
+ ("encodec", "EncodecModel"),
+ ("ernie", "ErnieModel"),
+ ("ernie4_5", "Ernie4_5Model"),
+ ("ernie4_5_moe", "Ernie4_5_MoeModel"),
+ ("ernie4_5_vl_moe", "Ernie4_5_VLMoeModel"),
+ ("esm", "EsmModel"),
+ ("eurobert", "EuroBertModel"),
+ ("evolla", "EvollaModel"),
+ ("exaone4", "Exaone4Model"),
+ ("exaone_moe", "ExaoneMoeModel"),
+ ("falcon", "FalconModel"),
+ ("falcon_h1", "FalconH1Model"),
+ ("falcon_mamba", "FalconMambaModel"),
+ ("fast_vlm", "FastVlmModel"),
+ ("fastspeech2_conformer", "FastSpeech2ConformerModel"),
+ ("fastspeech2_conformer_with_hifigan", "FastSpeech2ConformerWithHifiGan"),
+ ("flaubert", "FlaubertModel"),
+ ("flava", "FlavaModel"),
+ ("flex_olmo", "FlexOlmoModel"),
+ ("florence2", "Florence2Model"),
+ ("fnet", "FNetModel"),
+ ("focalnet", "FocalNetModel"),
+ ("fsmt", "FSMTModel"),
+ ("funnel", ("FunnelModel", "FunnelBaseModel")),
+ ("fuyu", "FuyuModel"),
+ ("gemma", "GemmaModel"),
+ ("gemma2", "Gemma2Model"),
+ ("gemma3", "Gemma3Model"),
+ ("gemma3_text", "Gemma3TextModel"),
+ ("gemma3n", "Gemma3nModel"),
+ ("gemma3n_audio", "Gemma3nAudioEncoder"),
+ ("gemma3n_text", "Gemma3nTextModel"),
+ ("gemma3n_vision", "TimmWrapperModel"),
+ ("gemma4", "Gemma4Model"),
+ ("gemma4_audio", "Gemma4AudioModel"),
+ ("gemma4_text", "Gemma4TextModel"),
+ ("gemma4_vision", "Gemma4VisionModel"),
+ ("git", "GitModel"),
+ ("glm", "GlmModel"),
+ ("glm4", "Glm4Model"),
+ ("glm46v", "Glm46VModel"),
+ ("glm4_moe", "Glm4MoeModel"),
+ ("glm4_moe_lite", "Glm4MoeLiteModel"),
+ ("glm4v", "Glm4vModel"),
+ ("glm4v_moe", "Glm4vMoeModel"),
+ ("glm4v_moe_text", "Glm4vMoeTextModel"),
+ ("glm4v_moe_vision", "Glm4vMoeVisionModel"),
+ ("glm4v_text", "Glm4vTextModel"),
+ ("glm4v_vision", "Glm4vVisionModel"),
+ ("glm_image", "GlmImageModel"),
+ ("glm_image_text", "GlmImageTextModel"),
+ ("glm_image_vision", "GlmImageVisionModel"),
+ ("glm_image_vqmodel", "GlmImageVQVAE"),
+ ("glm_moe_dsa", "GlmMoeDsaModel"),
+ ("glm_ocr", "GlmOcrModel"),
+ ("glm_ocr_text", "GlmOcrTextModel"),
+ ("glm_ocr_vision", "GlmOcrVisionModel"),
+ ("glmasr", "GlmAsrForConditionalGeneration"),
+ ("glmasr_encoder", "GlmAsrEncoder"),
+ ("glpn", "GLPNModel"),
+ ("got_ocr2", "GotOcr2Model"),
+ ("gpt-sw3", "GPT2Model"),
+ ("gpt2", "GPT2Model"),
+ ("gpt_bigcode", "GPTBigCodeModel"),
+ ("gpt_neo", "GPTNeoModel"),
+ ("gpt_neox", "GPTNeoXModel"),
+ ("gpt_neox_japanese", "GPTNeoXJapaneseModel"),
+ ("gpt_oss", "GptOssModel"),
+ ("gptj", "GPTJModel"),
+ ("granite", "GraniteModel"),
+ ("granitemoe", "GraniteMoeModel"),
+ ("granitemoehybrid", "GraniteMoeHybridModel"),
+ ("granitemoeshared", "GraniteMoeSharedModel"),
+ ("grounding-dino", "GroundingDinoModel"),
+ ("groupvit", "GroupViTModel"),
+ ("helium", "HeliumModel"),
+ ("hgnet_v2", "HGNetV2Backbone"),
+ ("hiera", "HieraModel"),
+ ("higgs_audio_v2", "HiggsAudioV2ForConditionalGeneration"),
+ ("higgs_audio_v2_tokenizer", "HiggsAudioV2TokenizerModel"),
+ ("hubert", "HubertModel"),
+ ("hunyuan_v1_dense", "HunYuanDenseV1Model"),
+ ("hunyuan_v1_moe", "HunYuanMoEV1Model"),
+ ("ibert", "IBertModel"),
+ ("idefics", "IdeficsModel"),
+ ("idefics2", "Idefics2Model"),
+ ("idefics3", "Idefics3Model"),
+ ("idefics3_vision", "Idefics3VisionTransformer"),
+ ("ijepa", "IJepaModel"),
+ ("imagegpt", "ImageGPTModel"),
+ ("informer", "InformerModel"),
+ ("instructblip", "InstructBlipModel"),
+ ("instructblipvideo", "InstructBlipVideoModel"),
+ ("internvl", "InternVLModel"),
+ ("internvl_vision", "InternVLVisionModel"),
+ ("jais2", "Jais2Model"),
+ ("jamba", "JambaModel"),
+ ("janus", "JanusModel"),
+ ("jetmoe", "JetMoeModel"),
+ ("jina_embeddings_v3", "JinaEmbeddingsV3Model"),
+ ("kosmos-2", "Kosmos2Model"),
+ ("kosmos-2.5", "Kosmos2_5Model"),
+ ("kyutai_speech_to_text", "KyutaiSpeechToTextModel"),
+ ("lasr_ctc", "LasrForCTC"),
+ ("lasr_encoder", "LasrEncoder"),
+ ("layoutlm", "LayoutLMModel"),
+ ("layoutlmv2", "LayoutLMv2Model"),
+ ("layoutlmv3", "LayoutLMv3Model"),
+ ("led", "LEDModel"),
+ ("levit", "LevitModel"),
+ ("lfm2", "Lfm2Model"),
+ ("lfm2_moe", "Lfm2MoeModel"),
+ ("lfm2_vl", "Lfm2VlModel"),
+ ("lightglue", "LightGlueForKeypointMatching"),
+ ("lighton_ocr", "LightOnOcrModel"),
+ ("lilt", "LiltModel"),
+ ("llama", "LlamaModel"),
+ ("llama4", "Llama4ForConditionalGeneration"),
+ ("llama4_text", "Llama4TextModel"),
+ ("llava", "LlavaModel"),
+ ("llava_next", "LlavaNextModel"),
+ ("llava_next_video", "LlavaNextVideoModel"),
+ ("llava_onevision", "LlavaOnevisionModel"),
+ ("longcat_flash", "LongcatFlashModel"),
+ ("longformer", "LongformerModel"),
+ ("longt5", "LongT5Model"),
+ ("luke", "LukeModel"),
+ ("lw_detr", "LwDetrModel"),
+ ("lxmert", "LxmertModel"),
+ ("m2m_100", "M2M100Model"),
+ ("mamba", "MambaModel"),
+ ("mamba2", "Mamba2Model"),
+ ("marian", "MarianModel"),
+ ("markuplm", "MarkupLMModel"),
+ ("mask2former", "Mask2FormerModel"),
+ ("maskformer", "MaskFormerModel"),
+ ("maskformer-swin", "MaskFormerSwinModel"),
+ ("mbart", "MBartModel"),
+ ("megatron-bert", "MegatronBertModel"),
+ ("metaclip_2", "MetaClip2Model"),
+ ("mgp-str", "MgpstrForSceneTextRecognition"),
+ ("mimi", "MimiModel"),
+ ("minimax", "MiniMaxModel"),
+ ("minimax_m2", "MiniMaxM2Model"),
+ ("ministral", "MinistralModel"),
+ ("ministral3", "Ministral3Model"),
+ ("mistral", "MistralModel"),
+ ("mistral3", "Mistral3Model"),
+ ("mistral4", "Mistral4Model"),
+ ("mixtral", "MixtralModel"),
+ ("mlcd", "MLCDVisionModel"), # Keep this to make some original hub repositories (from `DeepGlint-AI`) works
+ ("mlcd_vision_model", "MLCDVisionModel"),
+ ("mllama", "MllamaModel"),
+ ("mm-grounding-dino", "MMGroundingDinoModel"),
+ ("mobilebert", "MobileBertModel"),
+ ("mobilenet_v1", "MobileNetV1Model"),
+ ("mobilenet_v2", "MobileNetV2Model"),
+ ("mobilevit", "MobileViTModel"),
+ ("mobilevitv2", "MobileViTV2Model"),
+ ("modernbert", "ModernBertModel"),
+ ("modernbert-decoder", "ModernBertDecoderModel"),
+ ("modernvbert", "ModernVBertModel"),
+ ("moonshine", "MoonshineModel"),
+ ("moonshine_streaming", "MoonshineStreamingModel"),
+ ("moshi", "MoshiModel"),
+ ("mpnet", "MPNetModel"),
+ ("mpt", "MptModel"),
+ ("mra", "MraModel"),
+ ("mt5", "MT5Model"),
+ ("musicflamingo", "MusicFlamingoForConditionalGeneration"),
+ ("musicflamingo_encoder", "AudioFlamingo3Encoder"),
+ ("musicgen", "MusicgenModel"),
+ ("musicgen_melody", "MusicgenMelodyModel"),
+ ("mvp", "MvpModel"),
+ ("nanochat", "NanoChatModel"),
+ ("nemotron", "NemotronModel"),
+ ("nemotron_h", "NemotronHModel"),
+ ("nllb-moe", "NllbMoeModel"),
+ ("nomic_bert", "NomicBertModel"),
+ ("nystromformer", "NystromformerModel"),
+ ("olmo", "OlmoModel"),
+ ("olmo2", "Olmo2Model"),
+ ("olmo3", "Olmo3Model"),
+ ("olmo_hybrid", "OlmoHybridModel"),
+ ("olmoe", "OlmoeModel"),
+ ("omdet-turbo", "OmDetTurboForObjectDetection"),
+ ("oneformer", "OneFormerModel"),
+ ("openai-gpt", "OpenAIGPTModel"),
+ ("opt", "OPTModel"),
+ ("ovis2", "Ovis2Model"),
+ ("owlv2", "Owlv2Model"),
+ ("owlvit", "OwlViTModel"),
+ ("paligemma", "PaliGemmaModel"),
+ ("parakeet_ctc", "ParakeetForCTC"),
+ ("parakeet_encoder", "ParakeetEncoder"),
+ ("patchtsmixer", "PatchTSMixerModel"),
+ ("patchtst", "PatchTSTModel"),
+ ("pe_audio", "PeAudioModel"),
+ ("pe_audio_encoder", "PeAudioEncoder"),
+ ("pe_audio_video", "PeAudioVideoModel"),
+ ("pe_audio_video_encoder", "PeAudioVideoEncoder"),
+ ("pe_video", "PeVideoModel"),
+ ("pe_video_encoder", "PeVideoEncoder"),
+ ("pegasus", "PegasusModel"),
+ ("pegasus_x", "PegasusXModel"),
+ ("perceiver", "PerceiverModel"),
+ ("perception_lm", "PerceptionLMModel"),
+ ("persimmon", "PersimmonModel"),
+ ("phi", "PhiModel"),
+ ("phi3", "Phi3Model"),
+ ("phi4_multimodal", "Phi4MultimodalModel"),
+ ("phimoe", "PhimoeModel"),
+ ("pi0", "PI0Model"),
+ ("pixio", "PixioModel"),
+ ("pixtral", "PixtralVisionModel"),
+ ("plbart", "PLBartModel"),
+ ("poolformer", "PoolFormerModel"),
+ ("pp_doclayout_v3", "PPDocLayoutV3Model"),
+ ("pp_ocrv5_mobile_rec", "PPOCRV5MobileRecModel"),
+ ("pp_ocrv5_server_rec", "PPOCRV5ServerRecModel"),
+ ("prophetnet", "ProphetNetModel"),
+ ("pvt", "PvtModel"),
+ ("pvt_v2", "PvtV2Model"),
+ ("qwen2", "Qwen2Model"),
+ ("qwen2_5_vl", "Qwen2_5_VLModel"),
+ ("qwen2_5_vl_text", "Qwen2_5_VLTextModel"),
+ ("qwen2_audio_encoder", "Qwen2AudioEncoder"),
+ ("qwen2_moe", "Qwen2MoeModel"),
+ ("qwen2_vl", "Qwen2VLModel"),
+ ("qwen2_vl_text", "Qwen2VLTextModel"),
+ ("qwen3", "Qwen3Model"),
+ ("qwen3_5", "Qwen3_5Model"),
+ ("qwen3_5_moe", "Qwen3_5MoeModel"),
+ ("qwen3_5_moe_text", "Qwen3_5MoeTextModel"),
+ ("qwen3_5_text", "Qwen3_5TextModel"),
+ ("qwen3_moe", "Qwen3MoeModel"),
+ ("qwen3_next", "Qwen3NextModel"),
+ ("qwen3_vl", "Qwen3VLModel"),
+ ("qwen3_vl_moe", "Qwen3VLMoeModel"),
+ ("qwen3_vl_moe_text", "Qwen3VLMoeTextModel"),
+ ("qwen3_vl_text", "Qwen3VLTextModel"),
+ ("recurrent_gemma", "RecurrentGemmaModel"),
+ ("reformer", "ReformerModel"),
+ ("regnet", "RegNetModel"),
+ ("rembert", "RemBertModel"),
+ ("resnet", "ResNetModel"),
+ ("roberta", "RobertaModel"),
+ ("roberta-prelayernorm", "RobertaPreLayerNormModel"),
+ ("roc_bert", "RoCBertModel"),
+ ("roformer", "RoFormerModel"),
+ ("rt_detr", "RTDetrModel"),
+ ("rt_detr_v2", "RTDetrV2Model"),
+ ("rwkv", "RwkvModel"),
+ ("sam", "SamModel"),
+ ("sam2", "Sam2Model"),
+ ("sam2_hiera_det_model", "Sam2HieraDetModel"),
+ ("sam2_video", "Sam2VideoModel"),
+ ("sam2_vision_model", "Sam2VisionModel"),
+ ("sam3", "Sam3Model"),
+ ("sam3_tracker", "Sam3TrackerModel"),
+ ("sam3_tracker", "Sam3TrackerModel"),
+ ("sam3_tracker_video", "Sam3TrackerVideoModel"),
+ ("sam3_video", "Sam3VideoModel"),
+ ("sam3_vision_model", "Sam3VisionModel"),
+ ("sam3_vit_model", "Sam3ViTModel"),
+ ("sam_hq", "SamHQModel"),
+ ("sam_hq_vision_model", "SamHQVisionModel"),
+ ("sam_vision_model", "SamVisionModel"),
+ ("seamless_m4t", "SeamlessM4TModel"),
+ ("seamless_m4t_v2", "SeamlessM4Tv2Model"),
+ ("seed_oss", "SeedOssModel"),
+ ("segformer", "SegformerModel"),
+ ("seggpt", "SegGptModel"),
+ ("sew", "SEWModel"),
+ ("sew-d", "SEWDModel"),
+ ("siglip", "SiglipModel"),
+ ("siglip2", "Siglip2Model"),
+ ("siglip2_vision_model", "Siglip2VisionModel"),
+ ("siglip_vision_model", "SiglipVisionModel"),
+ ("smollm3", "SmolLM3Model"),
+ ("smolvlm", "SmolVLMModel"),
+ ("smolvlm_vision", "SmolVLMVisionTransformer"),
+ ("solar_open", "SolarOpenModel"),
+ ("speech_to_text", "Speech2TextModel"),
+ ("speecht5", "SpeechT5Model"),
+ ("splinter", "SplinterModel"),
+ ("squeezebert", "SqueezeBertModel"),
+ ("stablelm", "StableLmModel"),
+ ("starcoder2", "Starcoder2Model"),
+ ("swiftformer", "SwiftFormerModel"),
+ ("swin", "SwinModel"),
+ ("swin2sr", "Swin2SRModel"),
+ ("swinv2", "Swinv2Model"),
+ ("switch_transformers", "SwitchTransformersModel"),
+ ("t5", "T5Model"),
+ ("t5gemma", "T5GemmaModel"),
+ ("t5gemma2", "T5Gemma2Model"),
+ ("t5gemma2_encoder", "T5Gemma2Encoder"),
+ ("table-transformer", "TableTransformerModel"),
+ ("tapas", "TapasModel"),
+ ("textnet", "TextNetModel"),
+ ("time_series_transformer", "TimeSeriesTransformerModel"),
+ ("timesfm", "TimesFmModel"),
+ ("timesfm2_5", "TimesFm2_5Model"),
+ ("timesformer", "TimesformerModel"),
+ ("timm_backbone", "TimmBackbone"),
+ ("timm_wrapper", "TimmWrapperModel"),
+ ("tvp", "TvpModel"),
+ ("udop", "UdopModel"),
+ ("umt5", "UMT5Model"),
+ ("unispeech", "UniSpeechModel"),
+ ("unispeech-sat", "UniSpeechSatModel"),
+ ("univnet", "UnivNetModel"),
+ ("uvdoc", "UVDocModel"),
+ ("vaultgemma", "VaultGemmaModel"),
+ ("vibevoice_acoustic_tokenizer", "VibeVoiceAcousticTokenizerModel"),
+ ("vibevoice_acoustic_tokenizer_decoder", "VibeVoiceAcousticTokenizerDecoderModel"),
+ ("vibevoice_acoustic_tokenizer_encoder", "VibeVoiceAcousticTokenizerEncoderModel"),
+ ("vibevoice_asr", "VibeVoiceAsrForConditionalGeneration"),
+ ("video_llama_3", "VideoLlama3Model"),
+ ("video_llama_3_vision", "VideoLlama3VisionModel"),
+ ("video_llava", "VideoLlavaModel"),
+ ("videomae", "VideoMAEModel"),
+ ("vilt", "ViltModel"),
+ ("vipllava", "VipLlavaModel"),
+ ("vision-text-dual-encoder", "VisionTextDualEncoderModel"),
+ ("visual_bert", "VisualBertModel"),
+ ("vit", "ViTModel"),
+ ("vit_mae", "ViTMAEModel"),
+ ("vit_msn", "ViTMSNModel"),
+ ("vitdet", "VitDetModel"),
+ ("vits", "VitsModel"),
+ ("vivit", "VivitModel"),
+ ("vjepa2", "VJEPA2Model"),
+ ("voxtral", "VoxtralForConditionalGeneration"),
+ ("voxtral_encoder", "VoxtralEncoder"),
+ ("voxtral_realtime", "VoxtralRealtimeForConditionalGeneration"),
+ ("voxtral_realtime_encoder", "VoxtralRealtimeEncoder"),
+ ("voxtral_realtime_text", "VoxtralRealtimeTextModel"),
+ ("wav2vec2", "Wav2Vec2Model"),
+ ("wav2vec2-bert", "Wav2Vec2BertModel"),
+ ("wav2vec2-conformer", "Wav2Vec2ConformerModel"),
+ ("wavlm", "WavLMModel"),
+ ("whisper", "WhisperModel"),
+ ("xclip", "XCLIPModel"),
+ ("xcodec", "XcodecModel"),
+ ("xglm", "XGLMModel"),
+ ("xlm", "XLMModel"),
+ ("xlm-roberta", "XLMRobertaModel"),
+ ("xlm-roberta-xl", "XLMRobertaXLModel"),
+ ("xlnet", "XLNetModel"),
+ ("xlstm", "xLSTMModel"),
+ ("xmod", "XmodModel"),
+ ("yolos", "YolosModel"),
+ ("yoso", "YosoModel"),
+ ("youtu", "YoutuModel"),
+ ("zamba", "ZambaModel"),
+ ("zamba2", "Zamba2Model"),
+ ]
+)
+
+MODEL_FOR_PRETRAINING_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for pre-training mapping
+ ("albert", "AlbertForPreTraining"),
+ ("audioflamingo3", "AudioFlamingo3ForConditionalGeneration"),
+ ("bart", "BartForConditionalGeneration"),
+ ("bert", "BertForPreTraining"),
+ ("big_bird", "BigBirdForPreTraining"),
+ ("bloom", "BloomForCausalLM"),
+ ("camembert", "CamembertForMaskedLM"),
+ ("colmodernvbert", "ColModernVBertForRetrieval"),
+ ("colpali", "ColPaliForRetrieval"),
+ ("colqwen2", "ColQwen2ForRetrieval"),
+ ("ctrl", "CTRLLMHeadModel"),
+ ("data2vec-text", "Data2VecTextForMaskedLM"),
+ ("deberta", "DebertaForMaskedLM"),
+ ("deberta-v2", "DebertaV2ForMaskedLM"),
+ ("distilbert", "DistilBertForMaskedLM"),
+ ("electra", "ElectraForPreTraining"),
+ ("ernie", "ErnieForPreTraining"),
+ ("evolla", "EvollaForProteinText2Text"),
+ ("exaone4", "Exaone4ForCausalLM"),
+ ("exaone_moe", "ExaoneMoeForCausalLM"),
+ ("falcon_mamba", "FalconMambaForCausalLM"),
+ ("flaubert", "FlaubertWithLMHeadModel"),
+ ("flava", "FlavaForPreTraining"),
+ ("florence2", "Florence2ForConditionalGeneration"),
+ ("fnet", "FNetForPreTraining"),
+ ("fsmt", "FSMTForConditionalGeneration"),
+ ("funnel", "FunnelForPreTraining"),
+ ("gemma3", "Gemma3ForConditionalGeneration"),
+ ("gemma4", "Gemma4ForConditionalGeneration"),
+ ("glmasr", "GlmAsrForConditionalGeneration"),
+ ("gpt-sw3", "GPT2LMHeadModel"),
+ ("gpt2", "GPT2LMHeadModel"),
+ ("gpt_bigcode", "GPTBigCodeForCausalLM"),
+ ("hiera", "HieraForPreTraining"),
+ ("ibert", "IBertForMaskedLM"),
+ ("idefics", "IdeficsForVisionText2Text"),
+ ("idefics2", "Idefics2ForConditionalGeneration"),
+ ("idefics3", "Idefics3ForConditionalGeneration"),
+ ("janus", "JanusForConditionalGeneration"),
+ ("layoutlm", "LayoutLMForMaskedLM"),
+ ("llava", "LlavaForConditionalGeneration"),
+ ("llava_next", "LlavaNextForConditionalGeneration"),
+ ("llava_next_video", "LlavaNextVideoForConditionalGeneration"),
+ ("llava_onevision", "LlavaOnevisionForConditionalGeneration"),
+ ("longformer", "LongformerForMaskedLM"),
+ ("luke", "LukeForMaskedLM"),
+ ("lxmert", "LxmertForPreTraining"),
+ ("mamba", "MambaForCausalLM"),
+ ("mamba2", "Mamba2ForCausalLM"),
+ ("megatron-bert", "MegatronBertForPreTraining"),
+ ("mistral3", "Mistral3ForConditionalGeneration"),
+ ("mistral4", "Mistral4ForCausalLM"),
+ ("mllama", "MllamaForConditionalGeneration"),
+ ("mobilebert", "MobileBertForPreTraining"),
+ ("mpnet", "MPNetForMaskedLM"),
+ ("mpt", "MptForCausalLM"),
+ ("mra", "MraForMaskedLM"),
+ ("musicflamingo", "MusicFlamingoForConditionalGeneration"),
+ ("mvp", "MvpForConditionalGeneration"),
+ ("nanochat", "NanoChatForCausalLM"),
+ ("nllb-moe", "NllbMoeForConditionalGeneration"),
+ ("openai-gpt", "OpenAIGPTLMHeadModel"),
+ ("paligemma", "PaliGemmaForConditionalGeneration"),
+ ("qwen2_audio", "Qwen2AudioForConditionalGeneration"),
+ ("roberta", "RobertaForMaskedLM"),
+ ("roberta-prelayernorm", "RobertaPreLayerNormForMaskedLM"),
+ ("roc_bert", "RoCBertForPreTraining"),
+ ("rwkv", "RwkvForCausalLM"),
+ ("splinter", "SplinterForPreTraining"),
+ ("squeezebert", "SqueezeBertForMaskedLM"),
+ ("switch_transformers", "SwitchTransformersForConditionalGeneration"),
+ ("t5", "T5ForConditionalGeneration"),
+ ("t5gemma", "T5GemmaForConditionalGeneration"),
+ ("t5gemma2", "T5Gemma2ForConditionalGeneration"),
+ ("tapas", "TapasForMaskedLM"),
+ ("unispeech", "UniSpeechForPreTraining"),
+ ("unispeech-sat", "UniSpeechSatForPreTraining"),
+ ("vibevoice_asr", "VibeVoiceAsrForConditionalGeneration"),
+ ("video_llava", "VideoLlavaForConditionalGeneration"),
+ ("videomae", "VideoMAEForPreTraining"),
+ ("vipllava", "VipLlavaForConditionalGeneration"),
+ ("visual_bert", "VisualBertForPreTraining"),
+ ("vit_mae", "ViTMAEForPreTraining"),
+ ("voxtral", "VoxtralForConditionalGeneration"),
+ ("voxtral_realtime", "VoxtralRealtimeForConditionalGeneration"),
+ ("wav2vec2", "Wav2Vec2ForPreTraining"),
+ ("wav2vec2-conformer", "Wav2Vec2ConformerForPreTraining"),
+ ("xlm", "XLMWithLMHeadModel"),
+ ("xlm-roberta", "XLMRobertaForMaskedLM"),
+ ("xlm-roberta-xl", "XLMRobertaXLForMaskedLM"),
+ ("xlnet", "XLNetLMHeadModel"),
+ ("xlstm", "xLSTMForCausalLM"),
+ ("xmod", "XmodForMaskedLM"),
+ ]
+)
+
+MODEL_FOR_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Causal LM mapping
+ ("afmoe", "AfmoeForCausalLM"),
+ ("apertus", "ApertusForCausalLM"),
+ ("arcee", "ArceeForCausalLM"),
+ ("aria_text", "AriaTextForCausalLM"),
+ ("bamba", "BambaForCausalLM"),
+ ("bart", "BartForCausalLM"),
+ ("bert", "BertLMHeadModel"),
+ ("bert-generation", "BertGenerationDecoder"),
+ ("big_bird", "BigBirdForCausalLM"),
+ ("bigbird_pegasus", "BigBirdPegasusForCausalLM"),
+ ("biogpt", "BioGptForCausalLM"),
+ ("bitnet", "BitNetForCausalLM"),
+ ("blenderbot", "BlenderbotForCausalLM"),
+ ("blenderbot-small", "BlenderbotSmallForCausalLM"),
+ ("bloom", "BloomForCausalLM"),
+ ("blt", "BltForCausalLM"),
+ ("camembert", "CamembertForCausalLM"),
+ ("code_llama", "LlamaForCausalLM"),
+ ("codegen", "CodeGenForCausalLM"),
+ ("cohere", "CohereForCausalLM"),
+ ("cohere2", "Cohere2ForCausalLM"),
+ ("cpmant", "CpmAntForCausalLM"),
+ ("ctrl", "CTRLLMHeadModel"),
+ ("cwm", "CwmForCausalLM"),
+ ("data2vec-text", "Data2VecTextForCausalLM"),
+ ("dbrx", "DbrxForCausalLM"),
+ ("deepseek_v2", "DeepseekV2ForCausalLM"),
+ ("deepseek_v3", "DeepseekV3ForCausalLM"),
+ ("diffllama", "DiffLlamaForCausalLM"),
+ ("doge", "DogeForCausalLM"),
+ ("dots1", "Dots1ForCausalLM"),
+ ("electra", "ElectraForCausalLM"),
+ ("emu3", "Emu3ForCausalLM"),
+ ("ernie", "ErnieForCausalLM"),
+ ("ernie4_5", "Ernie4_5ForCausalLM"),
+ ("ernie4_5_moe", "Ernie4_5_MoeForCausalLM"),
+ ("exaone4", "Exaone4ForCausalLM"),
+ ("exaone_moe", "ExaoneMoeForCausalLM"),
+ ("falcon", "FalconForCausalLM"),
+ ("falcon_h1", "FalconH1ForCausalLM"),
+ ("falcon_mamba", "FalconMambaForCausalLM"),
+ ("flex_olmo", "FlexOlmoForCausalLM"),
+ ("fuyu", "FuyuForCausalLM"),
+ ("gemma", "GemmaForCausalLM"),
+ ("gemma2", "Gemma2ForCausalLM"),
+ ("gemma3", "Gemma3ForConditionalGeneration"),
+ ("gemma3_text", "Gemma3ForCausalLM"),
+ ("gemma3n", "Gemma3nForConditionalGeneration"),
+ ("gemma3n_text", "Gemma3nForCausalLM"),
+ ("gemma4", "Gemma4ForConditionalGeneration"),
+ ("gemma4_text", "Gemma4ForCausalLM"),
+ ("git", "GitForCausalLM"),
+ ("glm", "GlmForCausalLM"),
+ ("glm4", "Glm4ForCausalLM"),
+ ("glm4_moe", "Glm4MoeForCausalLM"),
+ ("glm4_moe_lite", "Glm4MoeLiteForCausalLM"),
+ ("glm_moe_dsa", "GlmMoeDsaForCausalLM"),
+ ("got_ocr2", "GotOcr2ForConditionalGeneration"),
+ ("gpt-sw3", "GPT2LMHeadModel"),
+ ("gpt2", "GPT2LMHeadModel"),
+ ("gpt_bigcode", "GPTBigCodeForCausalLM"),
+ ("gpt_neo", "GPTNeoForCausalLM"),
+ ("gpt_neox", "GPTNeoXForCausalLM"),
+ ("gpt_neox_japanese", "GPTNeoXJapaneseForCausalLM"),
+ ("gpt_oss", "GptOssForCausalLM"),
+ ("gptj", "GPTJForCausalLM"),
+ ("granite", "GraniteForCausalLM"),
+ ("granitemoe", "GraniteMoeForCausalLM"),
+ ("granitemoehybrid", "GraniteMoeHybridForCausalLM"),
+ ("granitemoeshared", "GraniteMoeSharedForCausalLM"),
+ ("helium", "HeliumForCausalLM"),
+ ("hunyuan_v1_dense", "HunYuanDenseV1ForCausalLM"),
+ ("hunyuan_v1_moe", "HunYuanMoEV1ForCausalLM"),
+ ("jais2", "Jais2ForCausalLM"),
+ ("jamba", "JambaForCausalLM"),
+ ("jetmoe", "JetMoeForCausalLM"),
+ ("lfm2", "Lfm2ForCausalLM"),
+ ("lfm2_moe", "Lfm2MoeForCausalLM"),
+ ("llama", "LlamaForCausalLM"),
+ ("llama4", "Llama4ForCausalLM"),
+ ("llama4_text", "Llama4ForCausalLM"),
+ ("longcat_flash", "LongcatFlashForCausalLM"),
+ ("mamba", "MambaForCausalLM"),
+ ("mamba2", "Mamba2ForCausalLM"),
+ ("marian", "MarianForCausalLM"),
+ ("mbart", "MBartForCausalLM"),
+ ("megatron-bert", "MegatronBertForCausalLM"),
+ ("minimax", "MiniMaxForCausalLM"),
+ ("minimax_m2", "MiniMaxM2ForCausalLM"),
+ ("ministral", "MinistralForCausalLM"),
+ ("ministral3", "Ministral3ForCausalLM"),
+ ("mistral", "MistralForCausalLM"),
+ ("mixtral", "MixtralForCausalLM"),
+ ("mllama", "MllamaForCausalLM"),
+ ("modernbert-decoder", "ModernBertDecoderForCausalLM"),
+ ("moshi", "MoshiForCausalLM"),
+ ("mpt", "MptForCausalLM"),
+ ("musicgen", "MusicgenForCausalLM"),
+ ("musicgen_melody", "MusicgenMelodyForCausalLM"),
+ ("mvp", "MvpForCausalLM"),
+ ("nanochat", "NanoChatForCausalLM"),
+ ("nemotron", "NemotronForCausalLM"),
+ ("nemotron_h", "NemotronHForCausalLM"),
+ ("olmo", "OlmoForCausalLM"),
+ ("olmo2", "Olmo2ForCausalLM"),
+ ("olmo3", "Olmo3ForCausalLM"),
+ ("olmo_hybrid", "OlmoHybridForCausalLM"),
+ ("olmoe", "OlmoeForCausalLM"),
+ ("openai-gpt", "OpenAIGPTLMHeadModel"),
+ ("opt", "OPTForCausalLM"),
+ ("pegasus", "PegasusForCausalLM"),
+ ("persimmon", "PersimmonForCausalLM"),
+ ("phi", "PhiForCausalLM"),
+ ("phi3", "Phi3ForCausalLM"),
+ ("phi4_multimodal", "Phi4MultimodalForCausalLM"),
+ ("phimoe", "PhimoeForCausalLM"),
+ ("plbart", "PLBartForCausalLM"),
+ ("prophetnet", "ProphetNetForCausalLM"),
+ ("qwen2", "Qwen2ForCausalLM"),
+ ("qwen2_moe", "Qwen2MoeForCausalLM"),
+ ("qwen3", "Qwen3ForCausalLM"),
+ ("qwen3_5", "Qwen3_5ForCausalLM"), # VLM compatibility
+ ("qwen3_5_moe", "Qwen3_5MoeForCausalLM"), # VLM compatibility
+ ("qwen3_5_moe_text", "Qwen3_5MoeForCausalLM"),
+ ("qwen3_5_text", "Qwen3_5ForCausalLM"),
+ ("qwen3_moe", "Qwen3MoeForCausalLM"),
+ ("qwen3_next", "Qwen3NextForCausalLM"),
+ ("recurrent_gemma", "RecurrentGemmaForCausalLM"),
+ ("reformer", "ReformerModelWithLMHead"),
+ ("rembert", "RemBertForCausalLM"),
+ ("roberta", "RobertaForCausalLM"),
+ ("roberta-prelayernorm", "RobertaPreLayerNormForCausalLM"),
+ ("roc_bert", "RoCBertForCausalLM"),
+ ("roformer", "RoFormerForCausalLM"),
+ ("rwkv", "RwkvForCausalLM"),
+ ("seed_oss", "SeedOssForCausalLM"),
+ ("smollm3", "SmolLM3ForCausalLM"),
+ ("solar_open", "SolarOpenForCausalLM"),
+ ("stablelm", "StableLmForCausalLM"),
+ ("starcoder2", "Starcoder2ForCausalLM"),
+ ("trocr", "TrOCRForCausalLM"),
+ ("vaultgemma", "VaultGemmaForCausalLM"),
+ ("whisper", "WhisperForCausalLM"),
+ ("xglm", "XGLMForCausalLM"),
+ ("xlm", "XLMWithLMHeadModel"),
+ ("xlm-roberta", "XLMRobertaForCausalLM"),
+ ("xlm-roberta-xl", "XLMRobertaXLForCausalLM"),
+ ("xlnet", "XLNetLMHeadModel"),
+ ("xlstm", "xLSTMForCausalLM"),
+ ("xmod", "XmodForCausalLM"),
+ ("youtu", "YoutuForCausalLM"),
+ ("zamba", "ZambaForCausalLM"),
+ ("zamba2", "Zamba2ForCausalLM"),
+ ]
+)
+
+MODEL_FOR_IMAGE_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Image mapping
+ ("aimv2_vision_model", "Aimv2VisionModel"),
+ ("beit", "BeitModel"),
+ ("bit", "BitModel"),
+ ("cohere2_vision", "Cohere2VisionModel"),
+ ("conditional_detr", "ConditionalDetrModel"),
+ ("convnext", "ConvNextModel"),
+ ("convnextv2", "ConvNextV2Model"),
+ ("dab-detr", "DabDetrModel"),
+ ("data2vec-vision", "Data2VecVisionModel"),
+ ("deformable_detr", "DeformableDetrModel"),
+ ("deit", "DeiTModel"),
+ ("depth_pro", "DepthProModel"),
+ ("detr", "DetrModel"),
+ ("dinat", "DinatModel"),
+ ("dinov2", "Dinov2Model"),
+ ("dinov2_with_registers", "Dinov2WithRegistersModel"),
+ ("dinov3_convnext", "DINOv3ConvNextModel"),
+ ("dinov3_vit", "DINOv3ViTModel"),
+ ("dpt", "DPTModel"),
+ ("efficientnet", "EfficientNetModel"),
+ ("focalnet", "FocalNetModel"),
+ ("glpn", "GLPNModel"),
+ ("hiera", "HieraModel"),
+ ("ijepa", "IJepaModel"),
+ ("imagegpt", "ImageGPTModel"),
+ ("levit", "LevitModel"),
+ ("llama4", "Llama4VisionModel"),
+ ("mlcd", "MLCDVisionModel"), # Keep this to make some original hub repositories (from `DeepGlint-AI`) works
+ ("mlcd_vision_model", "MLCDVisionModel"),
+ ("mllama", "MllamaVisionModel"),
+ ("mobilenet_v1", "MobileNetV1Model"),
+ ("mobilenet_v2", "MobileNetV2Model"),
+ ("mobilevit", "MobileViTModel"),
+ ("mobilevitv2", "MobileViTV2Model"),
+ ("pixio", "PixioModel"),
+ ("poolformer", "PoolFormerModel"),
+ ("pvt", "PvtModel"),
+ ("regnet", "RegNetModel"),
+ ("resnet", "ResNetModel"),
+ ("segformer", "SegformerModel"),
+ ("siglip_vision_model", "SiglipVisionModel"),
+ ("swiftformer", "SwiftFormerModel"),
+ ("swin", "SwinModel"),
+ ("swin2sr", "Swin2SRModel"),
+ ("swinv2", "Swinv2Model"),
+ ("table-transformer", "TableTransformerModel"),
+ ("timesformer", "TimesformerModel"),
+ ("timm_backbone", "TimmBackbone"),
+ ("timm_wrapper", "TimmWrapperModel"),
+ ("videomae", "VideoMAEModel"),
+ ("vit", "ViTModel"),
+ ("vit_mae", "ViTMAEModel"),
+ ("vit_msn", "ViTMSNModel"),
+ ("vitdet", "VitDetModel"),
+ ("vivit", "VivitModel"),
+ ("yolos", "YolosModel"),
+ ]
+)
+
+MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING_NAMES = OrderedDict(
+ [
+ ("deit", "DeiTForMaskedImageModeling"),
+ ("focalnet", "FocalNetForMaskedImageModeling"),
+ ("swin", "SwinForMaskedImageModeling"),
+ ("swinv2", "Swinv2ForMaskedImageModeling"),
+ ("vit", "ViTForMaskedImageModeling"),
+ ]
+)
+
+
+MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING_NAMES = OrderedDict(
+ # Model for Causal Image Modeling mapping
+ [
+ ("imagegpt", "ImageGPTForCausalImageModeling"),
+ ]
+)
+
+MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Image Classification mapping
+ ("beit", "BeitForImageClassification"),
+ ("bit", "BitForImageClassification"),
+ ("clip", "CLIPForImageClassification"),
+ ("convnext", "ConvNextForImageClassification"),
+ ("convnextv2", "ConvNextV2ForImageClassification"),
+ ("cvt", "CvtForImageClassification"),
+ ("data2vec-vision", "Data2VecVisionForImageClassification"),
+ (
+ "deit",
+ ("DeiTForImageClassification", "DeiTForImageClassificationWithTeacher"),
+ ),
+ ("dinat", "DinatForImageClassification"),
+ ("dinov2", "Dinov2ForImageClassification"),
+ ("dinov2_with_registers", "Dinov2WithRegistersForImageClassification"),
+ ("donut-swin", "DonutSwinForImageClassification"),
+ ("efficientnet", "EfficientNetForImageClassification"),
+ ("focalnet", "FocalNetForImageClassification"),
+ ("hgnet_v2", "HGNetV2ForImageClassification"),
+ ("hiera", "HieraForImageClassification"),
+ ("ijepa", "IJepaForImageClassification"),
+ ("imagegpt", "ImageGPTForImageClassification"),
+ (
+ "levit",
+ ("LevitForImageClassification", "LevitForImageClassificationWithTeacher"),
+ ),
+ ("metaclip_2", "MetaClip2ForImageClassification"),
+ ("mobilenet_v1", "MobileNetV1ForImageClassification"),
+ ("mobilenet_v2", "MobileNetV2ForImageClassification"),
+ ("mobilevit", "MobileViTForImageClassification"),
+ ("mobilevitv2", "MobileViTV2ForImageClassification"),
+ (
+ "perceiver",
+ (
+ "PerceiverForImageClassificationLearned",
+ "PerceiverForImageClassificationFourier",
+ "PerceiverForImageClassificationConvProcessing",
+ ),
+ ),
+ ("poolformer", "PoolFormerForImageClassification"),
+ ("pp_lcnet", "PPLCNetForImageClassification"),
+ ("pvt", "PvtForImageClassification"),
+ ("pvt_v2", "PvtV2ForImageClassification"),
+ ("regnet", "RegNetForImageClassification"),
+ ("resnet", "ResNetForImageClassification"),
+ ("segformer", "SegformerForImageClassification"),
+ ("shieldgemma2", "ShieldGemma2ForImageClassification"),
+ ("siglip", "SiglipForImageClassification"),
+ ("siglip2", "Siglip2ForImageClassification"),
+ ("swiftformer", "SwiftFormerForImageClassification"),
+ ("swin", "SwinForImageClassification"),
+ ("swinv2", "Swinv2ForImageClassification"),
+ ("textnet", "TextNetForImageClassification"),
+ ("timm_wrapper", "TimmWrapperForImageClassification"),
+ ("vit", "ViTForImageClassification"),
+ ("vit_msn", "ViTMSNForImageClassification"),
+ ]
+)
+
+MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES = OrderedDict(
+ [
+ # Do not add new models here, this class will be deprecated in the future.
+ # Model for Image Segmentation mapping
+ ("detr", "DetrForSegmentation"),
+ ]
+)
+
+MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Semantic Segmentation mapping
+ ("beit", "BeitForSemanticSegmentation"),
+ ("data2vec-vision", "Data2VecVisionForSemanticSegmentation"),
+ ("dpt", "DPTForSemanticSegmentation"),
+ ("mobilenet_v2", "MobileNetV2ForSemanticSegmentation"),
+ ("mobilevit", "MobileViTForSemanticSegmentation"),
+ ("mobilevitv2", "MobileViTV2ForSemanticSegmentation"),
+ ("segformer", "SegformerForSemanticSegmentation"),
+ ("upernet", "UperNetForSemanticSegmentation"),
+ ]
+)
+
+MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Instance Segmentation mapping
+ # MaskFormerForInstanceSegmentation can be removed from this mapping in v5
+ ("maskformer", "MaskFormerForInstanceSegmentation"),
+ ]
+)
+
+MODEL_FOR_UNIVERSAL_SEGMENTATION_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Universal Segmentation mapping
+ ("detr", "DetrForSegmentation"),
+ ("eomt", "EomtForUniversalSegmentation"),
+ ("eomt_dinov3", "EomtDinov3ForUniversalSegmentation"),
+ ("mask2former", "Mask2FormerForUniversalSegmentation"),
+ ("maskformer", "MaskFormerForInstanceSegmentation"),
+ ("oneformer", "OneFormerForUniversalSegmentation"),
+ ("videomt", "VideomtForUniversalSegmentation"),
+ ]
+)
+
+MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
+ [
+ ("timesformer", "TimesformerForVideoClassification"),
+ ("videomae", "VideoMAEForVideoClassification"),
+ ("vivit", "VivitForVideoClassification"),
+ ("vjepa2", "VJEPA2ForVideoClassification"),
+ ]
+)
+
+MODEL_FOR_RETRIEVAL_MAPPING_NAMES = OrderedDict(
+ [
+ ("colmodernvbert", "ColModernVBertForRetrieval"),
+ ("colpali", "ColPaliForRetrieval"),
+ ]
+)
+
+MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES = OrderedDict(
+ [
+ ("aria", "AriaForConditionalGeneration"),
+ ("aya_vision", "AyaVisionForConditionalGeneration"),
+ ("blip", "BlipForConditionalGeneration"),
+ ("blip-2", "Blip2ForConditionalGeneration"),
+ ("chameleon", "ChameleonForConditionalGeneration"),
+ ("cohere2_vision", "Cohere2VisionForConditionalGeneration"),
+ ("deepseek_vl", "DeepseekVLForConditionalGeneration"),
+ ("deepseek_vl_hybrid", "DeepseekVLHybridForConditionalGeneration"),
+ ("emu3", "Emu3ForConditionalGeneration"),
+ ("ernie4_5_vl_moe", "Ernie4_5_VLMoeForConditionalGeneration"),
+ ("evolla", "EvollaForProteinText2Text"),
+ ("fast_vlm", "FastVlmForConditionalGeneration"),
+ ("florence2", "Florence2ForConditionalGeneration"),
+ ("fuyu", "FuyuForCausalLM"),
+ ("gemma3", "Gemma3ForConditionalGeneration"),
+ ("gemma3n", "Gemma3nForConditionalGeneration"),
+ ("gemma4", "Gemma4ForConditionalGeneration"),
+ ("git", "GitForCausalLM"),
+ ("glm46v", "Glm46VForConditionalGeneration"),
+ ("glm4v", "Glm4vForConditionalGeneration"),
+ ("glm4v_moe", "Glm4vMoeForConditionalGeneration"),
+ ("glm_ocr", "GlmOcrForConditionalGeneration"),
+ ("got_ocr2", "GotOcr2ForConditionalGeneration"),
+ ("idefics", "IdeficsForVisionText2Text"),
+ ("idefics2", "Idefics2ForConditionalGeneration"),
+ ("idefics3", "Idefics3ForConditionalGeneration"),
+ ("instructblip", "InstructBlipForConditionalGeneration"),
+ ("instructblipvideo", "InstructBlipVideoForConditionalGeneration"),
+ ("internvl", "InternVLForConditionalGeneration"),
+ ("janus", "JanusForConditionalGeneration"),
+ ("kosmos-2", "Kosmos2ForConditionalGeneration"),
+ ("kosmos-2.5", "Kosmos2_5ForConditionalGeneration"),
+ ("lfm2_vl", "Lfm2VlForConditionalGeneration"),
+ ("lighton_ocr", "LightOnOcrForConditionalGeneration"),
+ ("llama4", "Llama4ForConditionalGeneration"),
+ ("llava", "LlavaForConditionalGeneration"),
+ ("llava_next", "LlavaNextForConditionalGeneration"),
+ ("llava_next_video", "LlavaNextVideoForConditionalGeneration"),
+ ("llava_onevision", "LlavaOnevisionForConditionalGeneration"),
+ ("mistral3", "Mistral3ForConditionalGeneration"),
+ ("mistral4", "Mistral4ForCausalLM"),
+ ("mllama", "MllamaForConditionalGeneration"),
+ ("ovis2", "Ovis2ForConditionalGeneration"),
+ ("paddleocr_vl", "PaddleOCRVLForConditionalGeneration"),
+ ("paligemma", "PaliGemmaForConditionalGeneration"),
+ ("perception_lm", "PerceptionLMForConditionalGeneration"),
+ ("pi0", "PI0ForConditionalGeneration"),
+ ("pix2struct", "Pix2StructForConditionalGeneration"),
+ ("pp_chart2table", "GotOcr2ForConditionalGeneration"),
+ ("qwen2_5_vl", "Qwen2_5_VLForConditionalGeneration"),
+ ("qwen2_vl", "Qwen2VLForConditionalGeneration"),
+ ("qwen3_5", "Qwen3_5ForConditionalGeneration"),
+ ("qwen3_5_moe", "Qwen3_5MoeForConditionalGeneration"),
+ ("qwen3_vl", "Qwen3VLForConditionalGeneration"),
+ ("qwen3_vl_moe", "Qwen3VLMoeForConditionalGeneration"),
+ ("shieldgemma2", "Gemma3ForConditionalGeneration"),
+ ("smolvlm", "SmolVLMForConditionalGeneration"),
+ ("t5gemma2", "T5Gemma2ForConditionalGeneration"),
+ ("udop", "UdopForConditionalGeneration"),
+ ("video_llama_3", "VideoLlama3ForConditionalGeneration"),
+ ("video_llava", "VideoLlavaForConditionalGeneration"),
+ ("vipllava", "VipLlavaForConditionalGeneration"),
+ ("vision-encoder-decoder", "VisionEncoderDecoderModel"),
+ ]
+)
+
+# Models that accept text and optionally multimodal data in inputs
+# and can generate text and optionally multimodal data.
+MODEL_FOR_MULTIMODAL_LM_MAPPING_NAMES = OrderedDict(
+ [
+ *list(MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES.items()),
+ ("glmasr", "GlmAsrForConditionalGeneration"),
+ ("granite_speech", "GraniteSpeechForConditionalGeneration"),
+ ("kyutai_speech_to_text", "KyutaiSpeechToTextForConditionalGeneration"),
+ ("phi4_multimodal", "Phi4MultimodalForCausalLM"),
+ ("qwen2_5_omni", "Qwen2_5OmniForConditionalGeneration"),
+ ("qwen2_audio", "Qwen2AudioForConditionalGeneration"),
+ ("qwen3_omni_moe", "Qwen3OmniMoeForConditionalGeneration"),
+ ("vibevoice_asr", "VibeVoiceAsrForConditionalGeneration"),
+ ("voxtral", "VoxtralForConditionalGeneration"),
+ ("voxtral_realtime", "VoxtralRealtimeForConditionalGeneration"),
+ ]
+)
+
+
+MODEL_FOR_MASKED_LM_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Masked LM mapping
+ ("albert", "AlbertForMaskedLM"),
+ ("bart", "BartForConditionalGeneration"),
+ ("bert", "BertForMaskedLM"),
+ ("big_bird", "BigBirdForMaskedLM"),
+ ("camembert", "CamembertForMaskedLM"),
+ ("convbert", "ConvBertForMaskedLM"),
+ ("data2vec-text", "Data2VecTextForMaskedLM"),
+ ("deberta", "DebertaForMaskedLM"),
+ ("deberta-v2", "DebertaV2ForMaskedLM"),
+ ("distilbert", "DistilBertForMaskedLM"),
+ ("electra", "ElectraForMaskedLM"),
+ ("ernie", "ErnieForMaskedLM"),
+ ("esm", "EsmForMaskedLM"),
+ ("eurobert", "EuroBertForMaskedLM"),
+ ("flaubert", "FlaubertWithLMHeadModel"),
+ ("fnet", "FNetForMaskedLM"),
+ ("funnel", "FunnelForMaskedLM"),
+ ("ibert", "IBertForMaskedLM"),
+ ("jina_embeddings_v3", "JinaEmbeddingsV3ForMaskedLM"),
+ ("layoutlm", "LayoutLMForMaskedLM"),
+ ("longformer", "LongformerForMaskedLM"),
+ ("luke", "LukeForMaskedLM"),
+ ("mbart", "MBartForConditionalGeneration"),
+ ("megatron-bert", "MegatronBertForMaskedLM"),
+ ("mobilebert", "MobileBertForMaskedLM"),
+ ("modernbert", "ModernBertForMaskedLM"),
+ ("modernvbert", "ModernVBertForMaskedLM"),
+ ("mpnet", "MPNetForMaskedLM"),
+ ("mra", "MraForMaskedLM"),
+ ("mvp", "MvpForConditionalGeneration"),
+ ("nomic_bert", "NomicBertForMaskedLM"),
+ ("nystromformer", "NystromformerForMaskedLM"),
+ ("perceiver", "PerceiverForMaskedLM"),
+ ("reformer", "ReformerForMaskedLM"),
+ ("rembert", "RemBertForMaskedLM"),
+ ("roberta", "RobertaForMaskedLM"),
+ ("roberta-prelayernorm", "RobertaPreLayerNormForMaskedLM"),
+ ("roc_bert", "RoCBertForMaskedLM"),
+ ("roformer", "RoFormerForMaskedLM"),
+ ("squeezebert", "SqueezeBertForMaskedLM"),
+ ("tapas", "TapasForMaskedLM"),
+ ("xlm", "XLMWithLMHeadModel"),
+ ("xlm-roberta", "XLMRobertaForMaskedLM"),
+ ("xlm-roberta-xl", "XLMRobertaXLForMaskedLM"),
+ ("xmod", "XmodForMaskedLM"),
+ ("yoso", "YosoForMaskedLM"),
+ ]
+)
+
+MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Object Detection mapping
+ ("conditional_detr", "ConditionalDetrForObjectDetection"),
+ ("d_fine", "DFineForObjectDetection"),
+ ("dab-detr", "DabDetrForObjectDetection"),
+ ("deformable_detr", "DeformableDetrForObjectDetection"),
+ ("detr", "DetrForObjectDetection"),
+ ("lw_detr", "LwDetrForObjectDetection"),
+ ("pp_doclayout_v2", "PPDocLayoutV2ForObjectDetection"),
+ ("pp_doclayout_v3", "PPDocLayoutV3ForObjectDetection"),
+ ("pp_ocrv5_mobile_det", "PPOCRV5MobileDetForObjectDetection"),
+ ("pp_ocrv5_server_det", "PPOCRV5ServerDetForObjectDetection"),
+ ("rt_detr", "RTDetrForObjectDetection"),
+ ("rt_detr_v2", "RTDetrV2ForObjectDetection"),
+ ("table-transformer", "TableTransformerForObjectDetection"),
+ ("yolos", "YolosForObjectDetection"),
+ ]
+)
+
+MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Zero Shot Object Detection mapping
+ ("grounding-dino", "GroundingDinoForObjectDetection"),
+ ("mm-grounding-dino", "MMGroundingDinoForObjectDetection"),
+ ("omdet-turbo", "OmDetTurboForObjectDetection"),
+ ("owlv2", "Owlv2ForObjectDetection"),
+ ("owlvit", "OwlViTForObjectDetection"),
+ ]
+)
+
+MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for depth estimation mapping
+ ("chmv2", "CHMv2ForDepthEstimation"),
+ ("depth_anything", "DepthAnythingForDepthEstimation"),
+ ("depth_pro", "DepthProForDepthEstimation"),
+ ("dpt", "DPTForDepthEstimation"),
+ ("glpn", "GLPNForDepthEstimation"),
+ ("prompt_depth_anything", "PromptDepthAnythingForDepthEstimation"),
+ ("zoedepth", "ZoeDepthForDepthEstimation"),
+ ]
+)
+
+
+MODEL_FOR_TEXT_RECOGNITION_MAPPING_NAMES = OrderedDict(
+ [
+ ("pp_ocrv5_mobile_rec", "PPOCRV5MobileRecForTextRecognition"),
+ ("pp_ocrv5_server_rec", "PPOCRV5ServerRecForTextRecognition"),
+ ]
+)
+
+
+MODEL_FOR_TABLE_RECOGNITION_MAPPING_NAMES = OrderedDict(
+ [
+ ("slanext", "SLANeXtForTableRecognition"),
+ ]
+)
+
+
+MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Seq2Seq Causal LM mapping
+ ("audioflamingo3", "AudioFlamingo3ForConditionalGeneration"),
+ ("bart", "BartForConditionalGeneration"),
+ ("bigbird_pegasus", "BigBirdPegasusForConditionalGeneration"),
+ ("blenderbot", "BlenderbotForConditionalGeneration"),
+ ("blenderbot-small", "BlenderbotSmallForConditionalGeneration"),
+ ("encoder-decoder", "EncoderDecoderModel"),
+ ("fsmt", "FSMTForConditionalGeneration"),
+ ("glmasr", "GlmAsrForConditionalGeneration"),
+ ("granite_speech", "GraniteSpeechForConditionalGeneration"),
+ ("led", "LEDForConditionalGeneration"),
+ ("longt5", "LongT5ForConditionalGeneration"),
+ ("m2m_100", "M2M100ForConditionalGeneration"),
+ ("marian", "MarianMTModel"),
+ ("mbart", "MBartForConditionalGeneration"),
+ ("mt5", "MT5ForConditionalGeneration"),
+ ("musicflamingo", "MusicFlamingoForConditionalGeneration"),
+ ("mvp", "MvpForConditionalGeneration"),
+ ("nllb-moe", "NllbMoeForConditionalGeneration"),
+ ("pegasus", "PegasusForConditionalGeneration"),
+ ("pegasus_x", "PegasusXForConditionalGeneration"),
+ ("plbart", "PLBartForConditionalGeneration"),
+ ("prophetnet", "ProphetNetForConditionalGeneration"),
+ ("qwen2_audio", "Qwen2AudioForConditionalGeneration"),
+ ("seamless_m4t", "SeamlessM4TForTextToText"),
+ ("seamless_m4t_v2", "SeamlessM4Tv2ForTextToText"),
+ ("switch_transformers", "SwitchTransformersForConditionalGeneration"),
+ ("t5", "T5ForConditionalGeneration"),
+ ("t5gemma", "T5GemmaForConditionalGeneration"),
+ ("t5gemma2", "T5Gemma2ForConditionalGeneration"),
+ ("umt5", "UMT5ForConditionalGeneration"),
+ ("vibevoice_asr", "VibeVoiceAsrForConditionalGeneration"),
+ ("voxtral", "VoxtralForConditionalGeneration"),
+ ("voxtral_realtime", "VoxtralRealtimeForConditionalGeneration"),
+ ]
+)
+
+
+MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES = OrderedDict(
+ [
+ ("cohere_asr", "CohereAsrForConditionalGeneration"),
+ ("dia", "DiaForConditionalGeneration"),
+ ("granite_speech", "GraniteSpeechForConditionalGeneration"),
+ ("kyutai_speech_to_text", "KyutaiSpeechToTextForConditionalGeneration"),
+ ("moonshine", "MoonshineForConditionalGeneration"),
+ ("moonshine_streaming", "MoonshineStreamingForConditionalGeneration"),
+ ("pop2piano", "Pop2PianoForConditionalGeneration"),
+ ("seamless_m4t", "SeamlessM4TForSpeechToText"),
+ ("seamless_m4t_v2", "SeamlessM4Tv2ForSpeechToText"),
+ ("speech-encoder-decoder", "SpeechEncoderDecoderModel"),
+ ("speech_to_text", "Speech2TextForConditionalGeneration"),
+ ("speecht5", "SpeechT5ForSpeechToText"),
+ ("vibevoice_asr", "VibeVoiceAsrForConditionalGeneration"),
+ ("voxtral", "VoxtralForConditionalGeneration"),
+ ("voxtral_realtime", "VoxtralRealtimeForConditionalGeneration"),
+ ("whisper", "WhisperForConditionalGeneration"),
+ ]
+)
+
+MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Sequence Classification mapping
+ ("albert", "AlbertForSequenceClassification"),
+ ("arcee", "ArceeForSequenceClassification"),
+ ("bart", "BartForSequenceClassification"),
+ ("bert", "BertForSequenceClassification"),
+ ("big_bird", "BigBirdForSequenceClassification"),
+ ("bigbird_pegasus", "BigBirdPegasusForSequenceClassification"),
+ ("biogpt", "BioGptForSequenceClassification"),
+ ("bloom", "BloomForSequenceClassification"),
+ ("camembert", "CamembertForSequenceClassification"),
+ ("canine", "CanineForSequenceClassification"),
+ ("code_llama", "LlamaForSequenceClassification"),
+ ("convbert", "ConvBertForSequenceClassification"),
+ ("ctrl", "CTRLForSequenceClassification"),
+ ("data2vec-text", "Data2VecTextForSequenceClassification"),
+ ("deberta", "DebertaForSequenceClassification"),
+ ("deberta-v2", "DebertaV2ForSequenceClassification"),
+ ("deepseek_v2", "DeepseekV2ForSequenceClassification"),
+ ("deepseek_v3", "DeepseekV3ForSequenceClassification"),
+ ("diffllama", "DiffLlamaForSequenceClassification"),
+ ("distilbert", "DistilBertForSequenceClassification"),
+ ("doge", "DogeForSequenceClassification"),
+ ("electra", "ElectraForSequenceClassification"),
+ ("ernie", "ErnieForSequenceClassification"),
+ ("esm", "EsmForSequenceClassification"),
+ ("eurobert", "EuroBertForSequenceClassification"),
+ ("exaone4", "Exaone4ForSequenceClassification"),
+ ("falcon", "FalconForSequenceClassification"),
+ ("flaubert", "FlaubertForSequenceClassification"),
+ ("fnet", "FNetForSequenceClassification"),
+ ("funnel", "FunnelForSequenceClassification"),
+ ("gemma", "GemmaForSequenceClassification"),
+ ("gemma2", "Gemma2ForSequenceClassification"),
+ ("gemma3", "Gemma3ForSequenceClassification"),
+ ("gemma3_text", "Gemma3TextForSequenceClassification"),
+ ("glm", "GlmForSequenceClassification"),
+ ("glm4", "Glm4ForSequenceClassification"),
+ ("gpt-sw3", "GPT2ForSequenceClassification"),
+ ("gpt2", "GPT2ForSequenceClassification"),
+ ("gpt_bigcode", "GPTBigCodeForSequenceClassification"),
+ ("gpt_neo", "GPTNeoForSequenceClassification"),
+ ("gpt_neox", "GPTNeoXForSequenceClassification"),
+ ("gpt_oss", "GptOssForSequenceClassification"),
+ ("gptj", "GPTJForSequenceClassification"),
+ ("helium", "HeliumForSequenceClassification"),
+ ("hunyuan_v1_dense", "HunYuanDenseV1ForSequenceClassification"),
+ ("hunyuan_v1_moe", "HunYuanMoEV1ForSequenceClassification"),
+ ("ibert", "IBertForSequenceClassification"),
+ ("jamba", "JambaForSequenceClassification"),
+ ("jetmoe", "JetMoeForSequenceClassification"),
+ ("jina_embeddings_v3", "JinaEmbeddingsV3ForSequenceClassification"),
+ ("layoutlm", "LayoutLMForSequenceClassification"),
+ ("layoutlmv2", "LayoutLMv2ForSequenceClassification"),
+ ("layoutlmv3", "LayoutLMv3ForSequenceClassification"),
+ ("lilt", "LiltForSequenceClassification"),
+ ("llama", "LlamaForSequenceClassification"),
+ ("longformer", "LongformerForSequenceClassification"),
+ ("luke", "LukeForSequenceClassification"),
+ ("markuplm", "MarkupLMForSequenceClassification"),
+ ("mbart", "MBartForSequenceClassification"),
+ ("megatron-bert", "MegatronBertForSequenceClassification"),
+ ("minimax", "MiniMaxForSequenceClassification"),
+ ("ministral", "MinistralForSequenceClassification"),
+ ("ministral3", "Ministral3ForSequenceClassification"),
+ ("mistral", "MistralForSequenceClassification"),
+ ("mistral4", "Mistral4ForSequenceClassification"),
+ ("mixtral", "MixtralForSequenceClassification"),
+ ("mobilebert", "MobileBertForSequenceClassification"),
+ ("modernbert", "ModernBertForSequenceClassification"),
+ ("modernbert-decoder", "ModernBertDecoderForSequenceClassification"),
+ ("modernvbert", "ModernVBertForSequenceClassification"),
+ ("mpnet", "MPNetForSequenceClassification"),
+ ("mpt", "MptForSequenceClassification"),
+ ("mra", "MraForSequenceClassification"),
+ ("mt5", "MT5ForSequenceClassification"),
+ ("mvp", "MvpForSequenceClassification"),
+ ("nemotron", "NemotronForSequenceClassification"),
+ ("nomic_bert", "NomicBertForSequenceClassification"),
+ ("nystromformer", "NystromformerForSequenceClassification"),
+ ("openai-gpt", "OpenAIGPTForSequenceClassification"),
+ ("opt", "OPTForSequenceClassification"),
+ ("perceiver", "PerceiverForSequenceClassification"),
+ ("persimmon", "PersimmonForSequenceClassification"),
+ ("phi", "PhiForSequenceClassification"),
+ ("phi3", "Phi3ForSequenceClassification"),
+ ("phimoe", "PhimoeForSequenceClassification"),
+ ("plbart", "PLBartForSequenceClassification"),
+ ("qwen2", "Qwen2ForSequenceClassification"),
+ ("qwen2_moe", "Qwen2MoeForSequenceClassification"),
+ ("qwen3", "Qwen3ForSequenceClassification"),
+ ("qwen3_5", "Qwen3_5ForSequenceClassification"),
+ ("qwen3_5_text", "Qwen3_5ForSequenceClassification"),
+ ("qwen3_moe", "Qwen3MoeForSequenceClassification"),
+ ("qwen3_next", "Qwen3NextForSequenceClassification"),
+ ("reformer", "ReformerForSequenceClassification"),
+ ("rembert", "RemBertForSequenceClassification"),
+ ("roberta", "RobertaForSequenceClassification"),
+ ("roberta-prelayernorm", "RobertaPreLayerNormForSequenceClassification"),
+ ("roc_bert", "RoCBertForSequenceClassification"),
+ ("roformer", "RoFormerForSequenceClassification"),
+ ("seed_oss", "SeedOssForSequenceClassification"),
+ ("smollm3", "SmolLM3ForSequenceClassification"),
+ ("squeezebert", "SqueezeBertForSequenceClassification"),
+ ("stablelm", "StableLmForSequenceClassification"),
+ ("starcoder2", "Starcoder2ForSequenceClassification"),
+ ("t5", "T5ForSequenceClassification"),
+ ("t5gemma", "T5GemmaForSequenceClassification"),
+ ("t5gemma2", "T5Gemma2ForSequenceClassification"),
+ ("tapas", "TapasForSequenceClassification"),
+ ("umt5", "UMT5ForSequenceClassification"),
+ ("xlm", "XLMForSequenceClassification"),
+ ("xlm-roberta", "XLMRobertaForSequenceClassification"),
+ ("xlm-roberta-xl", "XLMRobertaXLForSequenceClassification"),
+ ("xlnet", "XLNetForSequenceClassification"),
+ ("xmod", "XmodForSequenceClassification"),
+ ("yoso", "YosoForSequenceClassification"),
+ ("zamba", "ZambaForSequenceClassification"),
+ ("zamba2", "Zamba2ForSequenceClassification"),
+ ]
+)
+
+MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Question Answering mapping
+ ("albert", "AlbertForQuestionAnswering"),
+ ("arcee", "ArceeForQuestionAnswering"),
+ ("bart", "BartForQuestionAnswering"),
+ ("bert", "BertForQuestionAnswering"),
+ ("big_bird", "BigBirdForQuestionAnswering"),
+ ("bigbird_pegasus", "BigBirdPegasusForQuestionAnswering"),
+ ("bloom", "BloomForQuestionAnswering"),
+ ("camembert", "CamembertForQuestionAnswering"),
+ ("canine", "CanineForQuestionAnswering"),
+ ("convbert", "ConvBertForQuestionAnswering"),
+ ("data2vec-text", "Data2VecTextForQuestionAnswering"),
+ ("deberta", "DebertaForQuestionAnswering"),
+ ("deberta-v2", "DebertaV2ForQuestionAnswering"),
+ ("diffllama", "DiffLlamaForQuestionAnswering"),
+ ("distilbert", "DistilBertForQuestionAnswering"),
+ ("electra", "ElectraForQuestionAnswering"),
+ ("ernie", "ErnieForQuestionAnswering"),
+ ("exaone4", "Exaone4ForQuestionAnswering"),
+ ("falcon", "FalconForQuestionAnswering"),
+ ("flaubert", "FlaubertForQuestionAnsweringSimple"),
+ ("fnet", "FNetForQuestionAnswering"),
+ ("funnel", "FunnelForQuestionAnswering"),
+ ("gpt2", "GPT2ForQuestionAnswering"),
+ ("gpt_neo", "GPTNeoForQuestionAnswering"),
+ ("gpt_neox", "GPTNeoXForQuestionAnswering"),
+ ("gptj", "GPTJForQuestionAnswering"),
+ ("ibert", "IBertForQuestionAnswering"),
+ ("jina_embeddings_v3", "JinaEmbeddingsV3ForQuestionAnswering"),
+ ("layoutlmv2", "LayoutLMv2ForQuestionAnswering"),
+ ("layoutlmv3", "LayoutLMv3ForQuestionAnswering"),
+ ("led", "LEDForQuestionAnswering"),
+ ("lilt", "LiltForQuestionAnswering"),
+ ("llama", "LlamaForQuestionAnswering"),
+ ("longformer", "LongformerForQuestionAnswering"),
+ ("luke", "LukeForQuestionAnswering"),
+ ("lxmert", "LxmertForQuestionAnswering"),
+ ("markuplm", "MarkupLMForQuestionAnswering"),
+ ("mbart", "MBartForQuestionAnswering"),
+ ("megatron-bert", "MegatronBertForQuestionAnswering"),
+ ("minimax", "MiniMaxForQuestionAnswering"),
+ ("ministral", "MinistralForQuestionAnswering"),
+ ("ministral3", "Ministral3ForQuestionAnswering"),
+ ("mistral", "MistralForQuestionAnswering"),
+ ("mixtral", "MixtralForQuestionAnswering"),
+ ("mobilebert", "MobileBertForQuestionAnswering"),
+ ("modernbert", "ModernBertForQuestionAnswering"),
+ ("mpnet", "MPNetForQuestionAnswering"),
+ ("mpt", "MptForQuestionAnswering"),
+ ("mra", "MraForQuestionAnswering"),
+ ("mt5", "MT5ForQuestionAnswering"),
+ ("mvp", "MvpForQuestionAnswering"),
+ ("nemotron", "NemotronForQuestionAnswering"),
+ ("nystromformer", "NystromformerForQuestionAnswering"),
+ ("opt", "OPTForQuestionAnswering"),
+ ("qwen2", "Qwen2ForQuestionAnswering"),
+ ("qwen2_moe", "Qwen2MoeForQuestionAnswering"),
+ ("qwen3", "Qwen3ForQuestionAnswering"),
+ ("qwen3_moe", "Qwen3MoeForQuestionAnswering"),
+ ("qwen3_next", "Qwen3NextForQuestionAnswering"),
+ ("reformer", "ReformerForQuestionAnswering"),
+ ("rembert", "RemBertForQuestionAnswering"),
+ ("roberta", "RobertaForQuestionAnswering"),
+ ("roberta-prelayernorm", "RobertaPreLayerNormForQuestionAnswering"),
+ ("roc_bert", "RoCBertForQuestionAnswering"),
+ ("roformer", "RoFormerForQuestionAnswering"),
+ ("seed_oss", "SeedOssForQuestionAnswering"),
+ ("smollm3", "SmolLM3ForQuestionAnswering"),
+ ("splinter", "SplinterForQuestionAnswering"),
+ ("squeezebert", "SqueezeBertForQuestionAnswering"),
+ ("t5", "T5ForQuestionAnswering"),
+ ("umt5", "UMT5ForQuestionAnswering"),
+ ("xlm", "XLMForQuestionAnsweringSimple"),
+ ("xlm-roberta", "XLMRobertaForQuestionAnswering"),
+ ("xlm-roberta-xl", "XLMRobertaXLForQuestionAnswering"),
+ ("xlnet", "XLNetForQuestionAnsweringSimple"),
+ ("xmod", "XmodForQuestionAnswering"),
+ ("yoso", "YosoForQuestionAnswering"),
+ ]
+)
+
+MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Table Question Answering mapping
+ ("tapas", "TapasForQuestionAnswering"),
+ ]
+)
+
+MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
+ [
+ ("blip", "BlipForQuestionAnswering"),
+ ("blip-2", "Blip2ForConditionalGeneration"),
+ ("vilt", "ViltForQuestionAnswering"),
+ ]
+)
+
+MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES = OrderedDict(
+ [
+ ("layoutlm", "LayoutLMForQuestionAnswering"),
+ ("layoutlmv2", "LayoutLMv2ForQuestionAnswering"),
+ ("layoutlmv3", "LayoutLMv3ForQuestionAnswering"),
+ ]
+)
+
+MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Token Classification mapping
+ ("albert", "AlbertForTokenClassification"),
+ ("apertus", "ApertusForTokenClassification"),
+ ("arcee", "ArceeForTokenClassification"),
+ ("bert", "BertForTokenClassification"),
+ ("big_bird", "BigBirdForTokenClassification"),
+ ("biogpt", "BioGptForTokenClassification"),
+ ("bloom", "BloomForTokenClassification"),
+ ("bros", "BrosForTokenClassification"),
+ ("camembert", "CamembertForTokenClassification"),
+ ("canine", "CanineForTokenClassification"),
+ ("convbert", "ConvBertForTokenClassification"),
+ ("data2vec-text", "Data2VecTextForTokenClassification"),
+ ("deberta", "DebertaForTokenClassification"),
+ ("deberta-v2", "DebertaV2ForTokenClassification"),
+ ("deepseek_v3", "DeepseekV3ForTokenClassification"),
+ ("diffllama", "DiffLlamaForTokenClassification"),
+ ("distilbert", "DistilBertForTokenClassification"),
+ ("electra", "ElectraForTokenClassification"),
+ ("ernie", "ErnieForTokenClassification"),
+ ("esm", "EsmForTokenClassification"),
+ ("eurobert", "EuroBertForTokenClassification"),
+ ("exaone4", "Exaone4ForTokenClassification"),
+ ("falcon", "FalconForTokenClassification"),
+ ("flaubert", "FlaubertForTokenClassification"),
+ ("fnet", "FNetForTokenClassification"),
+ ("funnel", "FunnelForTokenClassification"),
+ ("gemma", "GemmaForTokenClassification"),
+ ("gemma2", "Gemma2ForTokenClassification"),
+ ("glm", "GlmForTokenClassification"),
+ ("glm4", "Glm4ForTokenClassification"),
+ ("gpt-sw3", "GPT2ForTokenClassification"),
+ ("gpt2", "GPT2ForTokenClassification"),
+ ("gpt_bigcode", "GPTBigCodeForTokenClassification"),
+ ("gpt_neo", "GPTNeoForTokenClassification"),
+ ("gpt_neox", "GPTNeoXForTokenClassification"),
+ ("gpt_oss", "GptOssForTokenClassification"),
+ ("helium", "HeliumForTokenClassification"),
+ ("ibert", "IBertForTokenClassification"),
+ ("jina_embeddings_v3", "JinaEmbeddingsV3ForTokenClassification"),
+ ("layoutlm", "LayoutLMForTokenClassification"),
+ ("layoutlmv2", "LayoutLMv2ForTokenClassification"),
+ ("layoutlmv3", "LayoutLMv3ForTokenClassification"),
+ ("lilt", "LiltForTokenClassification"),
+ ("llama", "LlamaForTokenClassification"),
+ ("longformer", "LongformerForTokenClassification"),
+ ("luke", "LukeForTokenClassification"),
+ ("markuplm", "MarkupLMForTokenClassification"),
+ ("megatron-bert", "MegatronBertForTokenClassification"),
+ ("minimax", "MiniMaxForTokenClassification"),
+ ("ministral", "MinistralForTokenClassification"),
+ ("ministral3", "Ministral3ForTokenClassification"),
+ ("mistral", "MistralForTokenClassification"),
+ ("mistral4", "Mistral4ForTokenClassification"),
+ ("mixtral", "MixtralForTokenClassification"),
+ ("mobilebert", "MobileBertForTokenClassification"),
+ ("modernbert", "ModernBertForTokenClassification"),
+ ("modernvbert", "ModernVBertForTokenClassification"),
+ ("mpnet", "MPNetForTokenClassification"),
+ ("mpt", "MptForTokenClassification"),
+ ("mra", "MraForTokenClassification"),
+ ("mt5", "MT5ForTokenClassification"),
+ ("nemotron", "NemotronForTokenClassification"),
+ ("nomic_bert", "NomicBertForTokenClassification"),
+ ("nystromformer", "NystromformerForTokenClassification"),
+ ("persimmon", "PersimmonForTokenClassification"),
+ ("phi", "PhiForTokenClassification"),
+ ("phi3", "Phi3ForTokenClassification"),
+ ("qwen2", "Qwen2ForTokenClassification"),
+ ("qwen2_moe", "Qwen2MoeForTokenClassification"),
+ ("qwen3", "Qwen3ForTokenClassification"),
+ ("qwen3_moe", "Qwen3MoeForTokenClassification"),
+ ("qwen3_next", "Qwen3NextForTokenClassification"),
+ ("rembert", "RemBertForTokenClassification"),
+ ("roberta", "RobertaForTokenClassification"),
+ ("roberta-prelayernorm", "RobertaPreLayerNormForTokenClassification"),
+ ("roc_bert", "RoCBertForTokenClassification"),
+ ("roformer", "RoFormerForTokenClassification"),
+ ("seed_oss", "SeedOssForTokenClassification"),
+ ("smollm3", "SmolLM3ForTokenClassification"),
+ ("squeezebert", "SqueezeBertForTokenClassification"),
+ ("stablelm", "StableLmForTokenClassification"),
+ ("starcoder2", "Starcoder2ForTokenClassification"),
+ ("t5", "T5ForTokenClassification"),
+ ("t5gemma", "T5GemmaForTokenClassification"),
+ ("t5gemma2", "T5Gemma2ForTokenClassification"),
+ ("umt5", "UMT5ForTokenClassification"),
+ ("xlm", "XLMForTokenClassification"),
+ ("xlm-roberta", "XLMRobertaForTokenClassification"),
+ ("xlm-roberta-xl", "XLMRobertaXLForTokenClassification"),
+ ("xlnet", "XLNetForTokenClassification"),
+ ("xmod", "XmodForTokenClassification"),
+ ("yoso", "YosoForTokenClassification"),
+ ]
+)
+
+MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Multiple Choice mapping
+ ("albert", "AlbertForMultipleChoice"),
+ ("bert", "BertForMultipleChoice"),
+ ("big_bird", "BigBirdForMultipleChoice"),
+ ("camembert", "CamembertForMultipleChoice"),
+ ("canine", "CanineForMultipleChoice"),
+ ("convbert", "ConvBertForMultipleChoice"),
+ ("data2vec-text", "Data2VecTextForMultipleChoice"),
+ ("deberta-v2", "DebertaV2ForMultipleChoice"),
+ ("distilbert", "DistilBertForMultipleChoice"),
+ ("electra", "ElectraForMultipleChoice"),
+ ("ernie", "ErnieForMultipleChoice"),
+ ("flaubert", "FlaubertForMultipleChoice"),
+ ("fnet", "FNetForMultipleChoice"),
+ ("funnel", "FunnelForMultipleChoice"),
+ ("ibert", "IBertForMultipleChoice"),
+ ("longformer", "LongformerForMultipleChoice"),
+ ("luke", "LukeForMultipleChoice"),
+ ("megatron-bert", "MegatronBertForMultipleChoice"),
+ ("mobilebert", "MobileBertForMultipleChoice"),
+ ("modernbert", "ModernBertForMultipleChoice"),
+ ("mpnet", "MPNetForMultipleChoice"),
+ ("mra", "MraForMultipleChoice"),
+ ("nystromformer", "NystromformerForMultipleChoice"),
+ ("rembert", "RemBertForMultipleChoice"),
+ ("roberta", "RobertaForMultipleChoice"),
+ ("roberta-prelayernorm", "RobertaPreLayerNormForMultipleChoice"),
+ ("roc_bert", "RoCBertForMultipleChoice"),
+ ("roformer", "RoFormerForMultipleChoice"),
+ ("squeezebert", "SqueezeBertForMultipleChoice"),
+ ("xlm", "XLMForMultipleChoice"),
+ ("xlm-roberta", "XLMRobertaForMultipleChoice"),
+ ("xlm-roberta-xl", "XLMRobertaXLForMultipleChoice"),
+ ("xlnet", "XLNetForMultipleChoice"),
+ ("xmod", "XmodForMultipleChoice"),
+ ("yoso", "YosoForMultipleChoice"),
+ ]
+)
+
+MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES = OrderedDict(
+ [
+ ("bert", "BertForNextSentencePrediction"),
+ ("ernie", "ErnieForNextSentencePrediction"),
+ ("fnet", "FNetForNextSentencePrediction"),
+ ("megatron-bert", "MegatronBertForNextSentencePrediction"),
+ ("mobilebert", "MobileBertForNextSentencePrediction"),
+ ]
+)
+
+MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Audio Classification mapping
+ ("audio-spectrogram-transformer", "ASTForAudioClassification"),
+ ("data2vec-audio", "Data2VecAudioForSequenceClassification"),
+ ("hubert", "HubertForSequenceClassification"),
+ ("sew", "SEWForSequenceClassification"),
+ ("sew-d", "SEWDForSequenceClassification"),
+ ("unispeech", "UniSpeechForSequenceClassification"),
+ ("unispeech-sat", "UniSpeechSatForSequenceClassification"),
+ ("wav2vec2", "Wav2Vec2ForSequenceClassification"),
+ ("wav2vec2-bert", "Wav2Vec2BertForSequenceClassification"),
+ ("wav2vec2-conformer", "Wav2Vec2ConformerForSequenceClassification"),
+ ("wavlm", "WavLMForSequenceClassification"),
+ ("whisper", "WhisperForAudioClassification"),
+ ]
+)
+
+MODEL_FOR_CTC_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Connectionist temporal classification (CTC) mapping
+ ("data2vec-audio", "Data2VecAudioForCTC"),
+ ("hubert", "HubertForCTC"),
+ ("lasr_ctc", "LasrForCTC"),
+ ("parakeet_ctc", "ParakeetForCTC"),
+ ("sew", "SEWForCTC"),
+ ("sew-d", "SEWDForCTC"),
+ ("unispeech", "UniSpeechForCTC"),
+ ("unispeech-sat", "UniSpeechSatForCTC"),
+ ("wav2vec2", "Wav2Vec2ForCTC"),
+ ("wav2vec2-bert", "Wav2Vec2BertForCTC"),
+ ("wav2vec2-conformer", "Wav2Vec2ConformerForCTC"),
+ ("wavlm", "WavLMForCTC"),
+ ]
+)
+
+MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Audio Classification mapping
+ ("data2vec-audio", "Data2VecAudioForAudioFrameClassification"),
+ ("unispeech-sat", "UniSpeechSatForAudioFrameClassification"),
+ ("wav2vec2", "Wav2Vec2ForAudioFrameClassification"),
+ ("wav2vec2-bert", "Wav2Vec2BertForAudioFrameClassification"),
+ ("wav2vec2-conformer", "Wav2Vec2ConformerForAudioFrameClassification"),
+ ("wavlm", "WavLMForAudioFrameClassification"),
+ ]
+)
+
+MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Audio Classification mapping
+ ("data2vec-audio", "Data2VecAudioForXVector"),
+ ("unispeech-sat", "UniSpeechSatForXVector"),
+ ("wav2vec2", "Wav2Vec2ForXVector"),
+ ("wav2vec2-bert", "Wav2Vec2BertForXVector"),
+ ("wav2vec2-conformer", "Wav2Vec2ConformerForXVector"),
+ ("wavlm", "WavLMForXVector"),
+ ]
+)
+
+MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Text-To-Spectrogram mapping
+ ("fastspeech2_conformer", "FastSpeech2ConformerModel"),
+ ("speecht5", "SpeechT5ForTextToSpeech"),
+ ]
+)
+
+MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Text-To-Waveform mapping
+ ("bark", "BarkModel"),
+ ("csm", "CsmForConditionalGeneration"),
+ ("fastspeech2_conformer_with_hifigan", "FastSpeech2ConformerWithHifiGan"),
+ ("higgs_audio_v2", "HiggsAudioV2ForConditionalGeneration"),
+ ("musicgen", "MusicgenForConditionalGeneration"),
+ ("musicgen_melody", "MusicgenMelodyForConditionalGeneration"),
+ ("qwen2_5_omni", "Qwen2_5OmniForConditionalGeneration"),
+ ("qwen3_omni_moe", "Qwen3OmniMoeForConditionalGeneration"),
+ ("seamless_m4t", "SeamlessM4TForTextToSpeech"),
+ ("seamless_m4t_v2", "SeamlessM4Tv2ForTextToSpeech"),
+ ("vits", "VitsModel"),
+ ]
+)
+
+MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
+ [
+ # Model for Zero Shot Image Classification mapping
+ ("align", "AlignModel"),
+ ("altclip", "AltCLIPModel"),
+ ("blip", "BlipModel"),
+ ("blip-2", "Blip2ForImageTextRetrieval"),
+ ("chinese_clip", "ChineseCLIPModel"),
+ ("clip", "CLIPModel"),
+ ("clipseg", "CLIPSegModel"),
+ ("metaclip_2", "MetaClip2Model"),
+ ("siglip", "SiglipModel"),
+ ("siglip2", "Siglip2Model"),
+ ]
+)
+
+MODEL_FOR_BACKBONE_MAPPING_NAMES = OrderedDict(
+ [
+ # Backbone mapping
+ ("beit", "BeitBackbone"),
+ ("bit", "BitBackbone"),
+ ("convnext", "ConvNextBackbone"),
+ ("convnextv2", "ConvNextV2Backbone"),
+ ("dinat", "DinatBackbone"),
+ ("dinov2", "Dinov2Backbone"),
+ ("dinov2_with_registers", "Dinov2WithRegistersBackbone"),
+ ("dinov3_convnext", "DINOv3ConvNextBackbone"),
+ ("dinov3_vit", "DINOv3ViTBackbone"),
+ ("focalnet", "FocalNetBackbone"),
+ ("hgnet_v2", "HGNetV2Backbone"),
+ ("hiera", "HieraBackbone"),
+ ("lw_detr_vit", "LwDetrViTBackbone"),
+ ("maskformer-swin", "MaskFormerSwinBackbone"),
+ ("pixio", "PixioBackbone"),
+ ("pp_lcnet", "PPLCNetBackbone"),
+ ("pp_lcnet_v3", "PPLCNetV3Backbone"),
+ ("pvt_v2", "PvtV2Backbone"),
+ ("resnet", "ResNetBackbone"),
+ ("rt_detr_resnet", "RTDetrResNetBackbone"),
+ ("swin", "SwinBackbone"),
+ ("swinv2", "Swinv2Backbone"),
+ ("textnet", "TextNetBackbone"),
+ ("timm_backbone", "TimmBackbone"),
+ ("uvdoc_backbone", "UVDocBackbone"),
+ ("vitdet", "VitDetBackbone"),
+ ("vitpose_backbone", "VitPoseBackbone"),
+ ]
+)
+
+MODEL_FOR_MASK_GENERATION_MAPPING_NAMES = OrderedDict(
+ [
+ ("edgetam", "EdgeTamModel"),
+ ("edgetam_video", "EdgeTamModel"),
+ ("sam", "SamModel"),
+ ("sam2", "Sam2Model"),
+ ("sam2_video", "Sam2Model"),
+ ("sam3_tracker", "Sam3TrackerModel"),
+ ("sam3_video", "Sam3TrackerModel"),
+ ("sam_hq", "SamHQModel"),
+ ]
+)
+
+
+MODEL_FOR_KEYPOINT_DETECTION_MAPPING_NAMES = OrderedDict(
+ [
+ ("superpoint", "SuperPointForKeypointDetection"),
+ ]
+)
+
+MODEL_FOR_KEYPOINT_MATCHING_MAPPING_NAMES = OrderedDict(
+ [
+ ("efficientloftr", "EfficientLoFTRForKeypointMatching"),
+ ("lightglue", "LightGlueForKeypointMatching"),
+ ("superglue", "SuperGlueForKeypointMatching"),
+ ]
+)
+
+MODEL_FOR_TEXT_ENCODING_MAPPING_NAMES = OrderedDict(
+ [
+ ("albert", "AlbertModel"),
+ ("bert", "BertModel"),
+ ("big_bird", "BigBirdModel"),
+ ("clip_text_model", "CLIPTextModel"),
+ ("data2vec-text", "Data2VecTextModel"),
+ ("deberta", "DebertaModel"),
+ ("deberta-v2", "DebertaV2Model"),
+ ("distilbert", "DistilBertModel"),
+ ("electra", "ElectraModel"),
+ ("emu3", "Emu3TextModel"),
+ ("flaubert", "FlaubertModel"),
+ ("ibert", "IBertModel"),
+ ("llama4", "Llama4TextModel"),
+ ("longformer", "LongformerModel"),
+ ("mllama", "MllamaTextModel"),
+ ("mobilebert", "MobileBertModel"),
+ ("mt5", "MT5EncoderModel"),
+ ("nystromformer", "NystromformerModel"),
+ ("reformer", "ReformerModel"),
+ ("rembert", "RemBertModel"),
+ ("roberta", "RobertaModel"),
+ ("roberta-prelayernorm", "RobertaPreLayerNormModel"),
+ ("roc_bert", "RoCBertModel"),
+ ("roformer", "RoFormerModel"),
+ ("squeezebert", "SqueezeBertModel"),
+ ("t5", "T5EncoderModel"),
+ ("t5gemma", "T5GemmaEncoderModel"),
+ ("umt5", "UMT5EncoderModel"),
+ ("xlm", "XLMModel"),
+ ("xlm-roberta", "XLMRobertaModel"),
+ ("xlm-roberta-xl", "XLMRobertaXLModel"),
+ ]
+)
+
+MODEL_FOR_TIME_SERIES_CLASSIFICATION_MAPPING_NAMES = OrderedDict(
+ [
+ ("patchtsmixer", "PatchTSMixerForTimeSeriesClassification"),
+ ("patchtst", "PatchTSTForClassification"),
+ ]
+)
+
+MODEL_FOR_TIME_SERIES_REGRESSION_MAPPING_NAMES = OrderedDict(
+ [
+ ("patchtsmixer", "PatchTSMixerForRegression"),
+ ("patchtst", "PatchTSTForRegression"),
+ ]
+)
+
+MODEL_FOR_TIME_SERIES_PREDICTION_MAPPING_NAMES = OrderedDict(
+ [
+ ("timesfm", "TimesFmModelForPrediction"),
+ ("timesfm2_5", "TimesFm2_5ModelForPrediction"),
+ ]
+)
+
+MODEL_FOR_IMAGE_TO_IMAGE_MAPPING_NAMES = OrderedDict(
+ [
+ ("swin2sr", "Swin2SRForImageSuperResolution"),
+ ]
+)
+
+MODEL_FOR_AUDIO_TOKENIZATION_NAMES = OrderedDict(
+ [
+ ("dac", "DacModel"),
+ ("higgs_audio_v2_tokenizer", "HiggsAudioV2TokenizerModel"),
+ ("vibevoice_acoustic_tokenizer", "VibeVoiceAcousticTokenizerModel"),
+ ]
+)
+
+MODEL_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_MAPPING_NAMES)
+MODEL_FOR_PRETRAINING_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_PRETRAINING_MAPPING_NAMES)
+MODEL_FOR_CAUSAL_LM_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_CAUSAL_LM_MAPPING_NAMES)
+MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING_NAMES
+)
+MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES
+)
+MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING_NAMES
+)
+MODEL_FOR_IMAGE_SEGMENTATION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_IMAGE_SEGMENTATION_MAPPING_NAMES
+)
+MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES
+)
+MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING_NAMES
+)
+MODEL_FOR_UNIVERSAL_SEGMENTATION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_UNIVERSAL_SEGMENTATION_MAPPING_NAMES
+)
+MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING_NAMES
+)
+MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING_NAMES
+)
+MODEL_FOR_MULTIMODAL_LM_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_MULTIMODAL_LM_MAPPING_NAMES)
+MODEL_FOR_RETRIEVAL_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_RETRIEVAL_MAPPING_NAMES)
+MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING_NAMES
+)
+MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES
+)
+MODEL_FOR_MASKED_LM_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_MASKED_LM_MAPPING_NAMES)
+MODEL_FOR_IMAGE_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_IMAGE_MAPPING_NAMES)
+MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING_NAMES
+)
+MODEL_FOR_OBJECT_DETECTION_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES)
+MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING_NAMES
+)
+MODEL_FOR_DEPTH_ESTIMATION_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES)
+MODEL_FOR_TEXT_RECOGNITION_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_TEXT_RECOGNITION_MAPPING_NAMES)
+MODEL_FOR_TABLE_RECOGNITION_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_TABLE_RECOGNITION_MAPPING_NAMES)
+MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES
+)
+MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES
+)
+MODEL_FOR_QUESTION_ANSWERING_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES
+)
+MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING_NAMES
+)
+MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES
+)
+MODEL_FOR_MULTIPLE_CHOICE_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES)
+MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING_NAMES
+)
+MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES
+)
+MODEL_FOR_CTC_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_CTC_MAPPING_NAMES)
+MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING_NAMES)
+MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING_NAMES
+)
+MODEL_FOR_AUDIO_XVECTOR_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_AUDIO_XVECTOR_MAPPING_NAMES)
+
+MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING_NAMES
+)
+
+MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING_NAMES)
+
+MODEL_FOR_BACKBONE_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_BACKBONE_MAPPING_NAMES)
+
+MODEL_FOR_MASK_GENERATION_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_MASK_GENERATION_MAPPING_NAMES)
+
+MODEL_FOR_KEYPOINT_DETECTION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_KEYPOINT_DETECTION_MAPPING_NAMES
+)
+
+MODEL_FOR_KEYPOINT_MATCHING_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_KEYPOINT_MATCHING_MAPPING_NAMES)
+
+MODEL_FOR_TEXT_ENCODING_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_TEXT_ENCODING_MAPPING_NAMES)
+
+MODEL_FOR_TIME_SERIES_CLASSIFICATION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_TIME_SERIES_CLASSIFICATION_MAPPING_NAMES
+)
+
+MODEL_FOR_TIME_SERIES_REGRESSION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_TIME_SERIES_REGRESSION_MAPPING_NAMES
+)
+
+MODEL_FOR_TIME_SERIES_PREDICTION_MAPPING = _LazyAutoMapping(
+ CONFIG_MAPPING_NAMES, MODEL_FOR_TIME_SERIES_PREDICTION_MAPPING_NAMES
+)
+
+MODEL_FOR_IMAGE_TO_IMAGE_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_IMAGE_TO_IMAGE_MAPPING_NAMES)
+
+MODEL_FOR_AUDIO_TOKENIZATION_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, MODEL_FOR_AUDIO_TOKENIZATION_NAMES)
+
+
+class AutoModelForMaskGeneration(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_MASK_GENERATION_MAPPING
+
+
+class AutoModelForKeypointDetection(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_KEYPOINT_DETECTION_MAPPING
+
+
+class AutoModelForKeypointMatching(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_KEYPOINT_MATCHING_MAPPING
+
+
+class AutoModelForTextEncoding(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_TEXT_ENCODING_MAPPING
+
+
+class AutoModelForImageToImage(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_IMAGE_TO_IMAGE_MAPPING
+
+
+class AutoModel(_BaseAutoModelClass):
+ _model_mapping = MODEL_MAPPING
+
+
+AutoModel = auto_class_update(AutoModel)
+
+
+class AutoModelForPreTraining(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_PRETRAINING_MAPPING
+
+
+AutoModelForPreTraining = auto_class_update(AutoModelForPreTraining, head_doc="pretraining")
+
+
+class AutoModelForCausalLM(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_CAUSAL_LM_MAPPING
+
+ # override to give better return typehint
+ @classmethod
+ def from_pretrained(
+ cls: type["AutoModelForCausalLM"],
+ pretrained_model_name_or_path: str | os.PathLike[str],
+ *model_args,
+ **kwargs,
+ ) -> "_BaseModelWithGenerate":
+ return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+
+
+AutoModelForCausalLM = auto_class_update(AutoModelForCausalLM, head_doc="causal language modeling")
+
+
+class AutoModelForMaskedLM(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_MASKED_LM_MAPPING
+
+
+AutoModelForMaskedLM = auto_class_update(AutoModelForMaskedLM, head_doc="masked language modeling")
+
+
+class AutoModelForSeq2SeqLM(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
+
+
+AutoModelForSeq2SeqLM = auto_class_update(
+ AutoModelForSeq2SeqLM,
+ head_doc="sequence-to-sequence language modeling",
+ checkpoint_for_example="google-t5/t5-base",
+)
+
+
+class AutoModelForSequenceClassification(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING
+
+
+AutoModelForSequenceClassification = auto_class_update(
+ AutoModelForSequenceClassification, head_doc="sequence classification"
+)
+
+
+class AutoModelForQuestionAnswering(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_QUESTION_ANSWERING_MAPPING
+
+
+AutoModelForQuestionAnswering = auto_class_update(AutoModelForQuestionAnswering, head_doc="question answering")
+
+
+class AutoModelForTableQuestionAnswering(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING
+
+
+AutoModelForTableQuestionAnswering = auto_class_update(
+ AutoModelForTableQuestionAnswering,
+ head_doc="table question answering",
+ checkpoint_for_example="google/tapas-base-finetuned-wtq",
+)
+
+
+class AutoModelForVisualQuestionAnswering(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING
+
+
+AutoModelForVisualQuestionAnswering = auto_class_update(
+ AutoModelForVisualQuestionAnswering,
+ head_doc="visual question answering",
+ checkpoint_for_example="dandelin/vilt-b32-finetuned-vqa",
+)
+
+
+class AutoModelForDocumentQuestionAnswering(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING
+
+
+AutoModelForDocumentQuestionAnswering = auto_class_update(
+ AutoModelForDocumentQuestionAnswering,
+ head_doc="document question answering",
+ checkpoint_for_example='impira/layoutlm-document-qa", revision="52e01b3',
+)
+
+
+class AutoModelForTokenClassification(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING
+
+
+AutoModelForTokenClassification = auto_class_update(AutoModelForTokenClassification, head_doc="token classification")
+
+
+class AutoModelForMultipleChoice(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_MULTIPLE_CHOICE_MAPPING
+
+
+AutoModelForMultipleChoice = auto_class_update(AutoModelForMultipleChoice, head_doc="multiple choice")
+
+
+class AutoModelForNextSentencePrediction(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING
+
+
+AutoModelForNextSentencePrediction = auto_class_update(
+ AutoModelForNextSentencePrediction, head_doc="next sentence prediction"
+)
+
+
+class AutoModelForImageClassification(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING
+
+
+AutoModelForImageClassification = auto_class_update(AutoModelForImageClassification, head_doc="image classification")
+
+
+class AutoModelForZeroShotImageClassification(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING
+
+
+AutoModelForZeroShotImageClassification = auto_class_update(
+ AutoModelForZeroShotImageClassification, head_doc="zero-shot image classification"
+)
+
+
+class AutoModelForImageSegmentation(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_IMAGE_SEGMENTATION_MAPPING
+
+
+AutoModelForImageSegmentation = auto_class_update(AutoModelForImageSegmentation, head_doc="image segmentation")
+
+
+class AutoModelForSemanticSegmentation(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING
+
+
+AutoModelForSemanticSegmentation = auto_class_update(
+ AutoModelForSemanticSegmentation, head_doc="semantic segmentation"
+)
+
+
+class AutoModelForTimeSeriesPrediction(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_TIME_SERIES_PREDICTION_MAPPING
+
+
+AutoModelForTimeSeriesPrediction = auto_class_update(
+ AutoModelForTimeSeriesPrediction, head_doc="time-series prediction"
+)
+
+
+class AutoModelForUniversalSegmentation(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_UNIVERSAL_SEGMENTATION_MAPPING
+
+
+AutoModelForUniversalSegmentation = auto_class_update(
+ AutoModelForUniversalSegmentation, head_doc="universal image segmentation"
+)
+
+
+class AutoModelForInstanceSegmentation(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING
+
+
+AutoModelForInstanceSegmentation = auto_class_update(
+ AutoModelForInstanceSegmentation, head_doc="instance segmentation"
+)
+
+
+class AutoModelForObjectDetection(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_OBJECT_DETECTION_MAPPING
+
+
+AutoModelForObjectDetection = auto_class_update(AutoModelForObjectDetection, head_doc="object detection")
+
+
+class AutoModelForZeroShotObjectDetection(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING
+
+
+AutoModelForZeroShotObjectDetection = auto_class_update(
+ AutoModelForZeroShotObjectDetection, head_doc="zero-shot object detection"
+)
+
+
+class AutoModelForDepthEstimation(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_DEPTH_ESTIMATION_MAPPING
+
+
+AutoModelForDepthEstimation = auto_class_update(AutoModelForDepthEstimation, head_doc="depth estimation")
+
+
+class AutoModelForTextRecognition(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_TEXT_RECOGNITION_MAPPING
+
+
+AutoModelForTextRecognition = auto_class_update(AutoModelForTextRecognition, head_doc="text recognition")
+
+
+class AutoModelForTableRecognition(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_TABLE_RECOGNITION_MAPPING
+
+
+AutoModelForTableRecognition = auto_class_update(AutoModelForTableRecognition, head_doc="table recognition")
+
+
+class AutoModelForVideoClassification(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING
+
+
+AutoModelForVideoClassification = auto_class_update(AutoModelForVideoClassification, head_doc="video classification")
+
+
+class AutoModelForImageTextToText(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING
+
+ # override to give better return typehint
+ @classmethod
+ def from_pretrained(
+ cls: type["AutoModelForImageTextToText"],
+ pretrained_model_name_or_path: str | os.PathLike[str],
+ *model_args,
+ **kwargs,
+ ) -> "_BaseModelWithGenerate":
+ return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+
+
+AutoModelForImageTextToText = auto_class_update(AutoModelForImageTextToText, head_doc="image-text-to-text modeling")
+
+
+class AutoModelForMultimodalLM(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_MULTIMODAL_LM_MAPPING
+
+
+AutoModelForMultimodalLM = auto_class_update(AutoModelForMultimodalLM, head_doc="multimodal generation")
+
+
+class AutoModelForAudioClassification(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING
+
+
+AutoModelForAudioClassification = auto_class_update(AutoModelForAudioClassification, head_doc="audio classification")
+
+
+class AutoModelForCTC(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_CTC_MAPPING
+
+
+AutoModelForCTC = auto_class_update(AutoModelForCTC, head_doc="connectionist temporal classification")
+
+
+class AutoModelForSpeechSeq2Seq(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING
+
+
+AutoModelForSpeechSeq2Seq = auto_class_update(
+ AutoModelForSpeechSeq2Seq, head_doc="sequence-to-sequence speech-to-text modeling"
+)
+
+
+class AutoModelForAudioFrameClassification(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING
+
+
+AutoModelForAudioFrameClassification = auto_class_update(
+ AutoModelForAudioFrameClassification, head_doc="audio frame (token) classification"
+)
+
+
+class AutoModelForAudioXVector(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_AUDIO_XVECTOR_MAPPING
+
+
+class AutoModelForTextToSpectrogram(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING
+
+
+class AutoModelForTextToWaveform(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING
+
+
+class AutoBackbone(_BaseAutoBackboneClass):
+ _model_mapping = MODEL_FOR_BACKBONE_MAPPING
+
+
+AutoModelForAudioXVector = auto_class_update(AutoModelForAudioXVector, head_doc="audio retrieval via x-vector")
+
+
+class AutoModelForMaskedImageModeling(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING
+
+
+AutoModelForMaskedImageModeling = auto_class_update(AutoModelForMaskedImageModeling, head_doc="masked image modeling")
+
+
+class AutoModelForAudioTokenization(_BaseAutoModelClass):
+ _model_mapping = MODEL_FOR_AUDIO_TOKENIZATION_MAPPING
+
+
+AutoModelForAudioTokenization = auto_class_update(
+ AutoModelForAudioTokenization, head_doc="audio tokenization through codebooks"
+)
+
+
+__all__ = [
+ "MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING",
+ "MODEL_FOR_AUDIO_FRAME_CLASSIFICATION_MAPPING",
+ "MODEL_FOR_AUDIO_TOKENIZATION_MAPPING",
+ "MODEL_FOR_AUDIO_XVECTOR_MAPPING",
+ "MODEL_FOR_BACKBONE_MAPPING",
+ "MODEL_FOR_CAUSAL_IMAGE_MODELING_MAPPING",
+ "MODEL_FOR_CAUSAL_LM_MAPPING",
+ "MODEL_FOR_CTC_MAPPING",
+ "MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING",
+ "MODEL_FOR_DEPTH_ESTIMATION_MAPPING",
+ "MODEL_FOR_TEXT_RECOGNITION_MAPPING",
+ "MODEL_FOR_TABLE_RECOGNITION_MAPPING",
+ "MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING",
+ "MODEL_FOR_IMAGE_MAPPING",
+ "MODEL_FOR_IMAGE_SEGMENTATION_MAPPING",
+ "MODEL_FOR_IMAGE_TO_IMAGE_MAPPING",
+ "MODEL_FOR_KEYPOINT_DETECTION_MAPPING",
+ "MODEL_FOR_KEYPOINT_MATCHING_MAPPING",
+ "MODEL_FOR_INSTANCE_SEGMENTATION_MAPPING",
+ "MODEL_FOR_MASKED_IMAGE_MODELING_MAPPING",
+ "MODEL_FOR_MASKED_LM_MAPPING",
+ "MODEL_FOR_MASK_GENERATION_MAPPING",
+ "MODEL_FOR_MULTIPLE_CHOICE_MAPPING",
+ "MODEL_FOR_NEXT_SENTENCE_PREDICTION_MAPPING",
+ "MODEL_FOR_OBJECT_DETECTION_MAPPING",
+ "MODEL_FOR_PRETRAINING_MAPPING",
+ "MODEL_FOR_QUESTION_ANSWERING_MAPPING",
+ "MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING",
+ "MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING",
+ "MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING",
+ "MODEL_FOR_SPEECH_SEQ_2_SEQ_MAPPING",
+ "MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING",
+ "MODEL_FOR_TEXT_ENCODING_MAPPING",
+ "MODEL_FOR_TEXT_TO_WAVEFORM_MAPPING",
+ "MODEL_FOR_TEXT_TO_SPECTROGRAM_MAPPING",
+ "MODEL_FOR_TIME_SERIES_PREDICTION_MAPPING",
+ "MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING",
+ "MODEL_FOR_UNIVERSAL_SEGMENTATION_MAPPING",
+ "MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING",
+ "MODEL_FOR_RETRIEVAL_MAPPING",
+ "MODEL_FOR_IMAGE_TEXT_TO_TEXT_MAPPING",
+ "MODEL_FOR_MULTIMODAL_LM_MAPPING",
+ "MODEL_FOR_VISUAL_QUESTION_ANSWERING_MAPPING",
+ "MODEL_MAPPING",
+ "MODEL_FOR_ZERO_SHOT_IMAGE_CLASSIFICATION_MAPPING",
+ "MODEL_FOR_ZERO_SHOT_OBJECT_DETECTION_MAPPING",
+ "MODEL_FOR_TIME_SERIES_CLASSIFICATION_MAPPING",
+ "MODEL_FOR_TIME_SERIES_REGRESSION_MAPPING",
+ "AutoModel",
+ "AutoBackbone",
+ "AutoModelForAudioClassification",
+ "AutoModelForAudioFrameClassification",
+ "AutoModelForAudioTokenization",
+ "AutoModelForAudioXVector",
+ "AutoModelForCausalLM",
+ "AutoModelForCTC",
+ "AutoModelForDepthEstimation",
+ "AutoModelForTextRecognition",
+ "AutoModelForTableRecognition",
+ "AutoModelForImageClassification",
+ "AutoModelForImageSegmentation",
+ "AutoModelForImageToImage",
+ "AutoModelForInstanceSegmentation",
+ "AutoModelForKeypointDetection",
+ "AutoModelForKeypointMatching",
+ "AutoModelForMaskGeneration",
+ "AutoModelForTextEncoding",
+ "AutoModelForMaskedImageModeling",
+ "AutoModelForMaskedLM",
+ "AutoModelForMultipleChoice",
+ "AutoModelForMultimodalLM",
+ "AutoModelForNextSentencePrediction",
+ "AutoModelForObjectDetection",
+ "AutoModelForPreTraining",
+ "AutoModelForQuestionAnswering",
+ "AutoModelForSemanticSegmentation",
+ "AutoModelForSeq2SeqLM",
+ "AutoModelForSequenceClassification",
+ "AutoModelForSpeechSeq2Seq",
+ "AutoModelForTableQuestionAnswering",
+ "AutoModelForTextToSpectrogram",
+ "AutoModelForTextToWaveform",
+ "AutoModelForTimeSeriesPrediction",
+ "AutoModelForTokenClassification",
+ "AutoModelForUniversalSegmentation",
+ "AutoModelForVideoClassification",
+ "AutoModelForVisualQuestionAnswering",
+ "AutoModelForDocumentQuestionAnswering",
+ "AutoModelForZeroShotImageClassification",
+ "AutoModelForZeroShotObjectDetection",
+ "AutoModelForImageTextToText",
+]
diff --git a/third_party/transformers/src/transformers/models/auto/processing_auto.py b/third_party/transformers/src/transformers/models/auto/processing_auto.py
new file mode 100644
index 0000000000000000000000000000000000000000..262480b71485e1714a189e12f11057e0882314ba
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/auto/processing_auto.py
@@ -0,0 +1,465 @@
+# Copyright 2021 The HuggingFace Inc. team.
+#
+# 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.
+"""AutoProcessor class."""
+
+import importlib
+import json
+from collections import OrderedDict
+from typing import TYPE_CHECKING
+
+# Build the list of all feature extractors
+from ...configuration_utils import PreTrainedConfig
+from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
+from ...feature_extraction_utils import FeatureExtractionMixin
+from ...image_processing_utils import ImageProcessingMixin
+from ...processing_utils import ProcessorMixin
+from ...tokenization_python import TOKENIZER_CONFIG_FILE
+from ...utils import FEATURE_EXTRACTOR_NAME, PROCESSOR_NAME, VIDEO_PROCESSOR_NAME, cached_file, logging
+from ...video_processing_utils import BaseVideoProcessor
+from .auto_factory import _LazyAutoMapping
+from .configuration_auto import (
+ CONFIG_MAPPING_NAMES,
+ AutoConfig,
+ model_type_to_module_name,
+ replace_list_option_in_docstrings,
+)
+from .feature_extraction_auto import AutoFeatureExtractor
+from .image_processing_auto import AutoImageProcessor
+from .tokenization_auto import AutoTokenizer
+from .video_processing_auto import AutoVideoProcessor
+
+
+logger = logging.get_logger(__name__)
+if TYPE_CHECKING:
+ # This significantly improves completion suggestion performance when
+ # the transformers package is used with Microsoft's Pylance language server.
+ PROCESSOR_MAPPING_NAMES: OrderedDict[str, str | None] = OrderedDict()
+else:
+ PROCESSOR_MAPPING_NAMES = OrderedDict(
+ [
+ ("aimv2", "CLIPProcessor"),
+ ("align", "AlignProcessor"),
+ ("altclip", "AltCLIPProcessor"),
+ ("aria", "AriaProcessor"),
+ ("audioflamingo3", "AudioFlamingo3Processor"),
+ ("aya_vision", "AyaVisionProcessor"),
+ ("bark", "BarkProcessor"),
+ ("blip", "BlipProcessor"),
+ ("blip-2", "Blip2Processor"),
+ ("bridgetower", "BridgeTowerProcessor"),
+ ("chameleon", "ChameleonProcessor"),
+ ("chinese_clip", "ChineseCLIPProcessor"),
+ ("clap", "ClapProcessor"),
+ ("clip", "CLIPProcessor"),
+ ("clipseg", "CLIPSegProcessor"),
+ ("clvp", "ClvpProcessor"),
+ ("cohere2_vision", "Cohere2VisionProcessor"),
+ ("cohere_asr", "CohereAsrProcessor"),
+ ("colmodernvbert", "ColModernVBertProcessor"),
+ ("colpali", "ColPaliProcessor"),
+ ("colqwen2", "ColQwen2Processor"),
+ ("deepseek_vl", "DeepseekVLProcessor"),
+ ("deepseek_vl_hybrid", "DeepseekVLHybridProcessor"),
+ ("dia", "DiaProcessor"),
+ ("edgetam", "Sam2Processor"),
+ ("emu3", "Emu3Processor"),
+ ("ernie4_5_vl_moe", "Ernie4_5_VLMoeProcessor"),
+ ("evolla", "EvollaProcessor"),
+ ("flava", "FlavaProcessor"),
+ ("florence2", "Florence2Processor"),
+ ("fuyu", "FuyuProcessor"),
+ ("gemma3", "Gemma3Processor"),
+ ("gemma3n", "Gemma3nProcessor"),
+ ("gemma4", "Gemma4Processor"),
+ ("git", "GitProcessor"),
+ ("glm46v", "Glm46VProcessor"),
+ ("glm4v", "Glm4vProcessor"),
+ ("glm4v_moe", "Glm4vProcessor"),
+ ("glm_image", "Glm4vProcessor"),
+ ("glmasr", "GlmAsrProcessor"),
+ ("got_ocr2", "GotOcr2Processor"),
+ ("granite_speech", "GraniteSpeechProcessor"),
+ ("grounding-dino", "GroundingDinoProcessor"),
+ ("groupvit", "CLIPProcessor"),
+ ("higgs_audio_v2", "HiggsAudioV2Processor"),
+ ("hubert", "Wav2Vec2Processor"),
+ ("idefics", "IdeficsProcessor"),
+ ("idefics2", "Idefics2Processor"),
+ ("idefics3", "Idefics3Processor"),
+ ("instructblip", "InstructBlipProcessor"),
+ ("instructblipvideo", "InstructBlipVideoProcessor"),
+ ("internvl", "InternVLProcessor"),
+ ("janus", "JanusProcessor"),
+ ("kosmos-2", "Kosmos2Processor"),
+ ("kosmos-2.5", "Kosmos2_5Processor"),
+ ("kyutai_speech_to_text", "KyutaiSpeechToTextProcessor"),
+ ("lasr_ctc", "LasrProcessor"),
+ ("lasr_encoder", "LasrProcessor"),
+ ("layoutlmv2", "LayoutLMv2Processor"),
+ ("layoutlmv3", "LayoutLMv3Processor"),
+ ("layoutxlm", "LayoutXLMProcessor"),
+ ("lfm2_vl", "Lfm2VlProcessor"),
+ ("lighton_ocr", "LightOnOcrProcessor"),
+ ("llama4", "Llama4Processor"),
+ ("llava", "LlavaProcessor"),
+ ("llava_next", "LlavaNextProcessor"),
+ ("llava_next_video", "LlavaNextVideoProcessor"),
+ ("llava_onevision", "LlavaOnevisionProcessor"),
+ ("markuplm", "MarkupLMProcessor"),
+ ("metaclip_2", "CLIPProcessor"),
+ ("mgp-str", "MgpstrProcessor"),
+ ("mistral3", "PixtralProcessor"),
+ ("mllama", "MllamaProcessor"),
+ ("mm-grounding-dino", "GroundingDinoProcessor"),
+ ("modernvbert", "Idefics3Processor"),
+ ("moonshine", "Wav2Vec2Processor"),
+ ("moonshine_streaming", "MoonshineStreamingProcessor"),
+ ("musicflamingo", "MusicFlamingoProcessor"),
+ ("omdet-turbo", "OmDetTurboProcessor"),
+ ("oneformer", "OneFormerProcessor"),
+ ("ovis2", "Ovis2Processor"),
+ ("owlv2", "Owlv2Processor"),
+ ("owlvit", "OwlViTProcessor"),
+ ("paddleocr_vl", "PaddleOCRVLProcessor"),
+ ("paligemma", "PaliGemmaProcessor"),
+ ("perception_lm", "PerceptionLMProcessor"),
+ ("phi4_multimodal", "Phi4MultimodalProcessor"),
+ ("pi0", "PI0Processor"),
+ ("pix2struct", "Pix2StructProcessor"),
+ ("pixtral", "PixtralProcessor"),
+ ("pop2piano", "Pop2PianoProcessor"),
+ ("pp_chart2table", "PPChart2TableProcessor"),
+ ("qwen2_5_omni", "Qwen2_5OmniProcessor"),
+ ("qwen2_5_vl", "Qwen2_5_VLProcessor"),
+ ("qwen2_audio", "Qwen2AudioProcessor"),
+ ("qwen2_vl", "Qwen2VLProcessor"),
+ ("qwen3_5", "Qwen3VLProcessor"),
+ ("qwen3_5_moe", "Qwen3VLProcessor"),
+ ("qwen3_omni_moe", "Qwen3OmniMoeProcessor"),
+ ("qwen3_vl", "Qwen3VLProcessor"),
+ ("qwen3_vl_moe", "Qwen3VLProcessor"),
+ ("sam", "SamProcessor"),
+ ("sam2", "Sam2Processor"),
+ ("sam3", "Sam3Processor"),
+ ("sam_hq", "SamHQProcessor"),
+ ("seamless_m4t", "SeamlessM4TProcessor"),
+ ("sew", "Wav2Vec2Processor"),
+ ("sew-d", "Wav2Vec2Processor"),
+ ("shieldgemma2", "ShieldGemma2Processor"),
+ ("siglip", "SiglipProcessor"),
+ ("siglip2", "Siglip2Processor"),
+ ("smolvlm", "SmolVLMProcessor"),
+ ("speech_to_text", "Speech2TextProcessor"),
+ ("speecht5", "SpeechT5Processor"),
+ ("t5gemma2", "Gemma3Processor"),
+ ("t5gemma2_encoder", "Gemma3Processor"),
+ ("trocr", "TrOCRProcessor"),
+ ("tvp", "TvpProcessor"),
+ ("udop", "UdopProcessor"),
+ ("unispeech", "Wav2Vec2Processor"),
+ ("unispeech-sat", "Wav2Vec2Processor"),
+ ("vibevoice_asr", "VibeVoiceAsrProcessor"),
+ ("video_llava", "VideoLlavaProcessor"),
+ ("vilt", "ViltProcessor"),
+ ("vipllava", "LlavaProcessor"),
+ ("vision-text-dual-encoder", "VisionTextDualEncoderProcessor"),
+ ("voxtral", "VoxtralProcessor"),
+ ("voxtral_realtime", "VoxtralRealtimeProcessor"),
+ ("wav2vec2", "Wav2Vec2Processor"),
+ ("wav2vec2-bert", "Wav2Vec2Processor"),
+ ("wav2vec2-conformer", "Wav2Vec2Processor"),
+ ("wavlm", "Wav2Vec2Processor"),
+ ("whisper", "WhisperProcessor"),
+ ("xclip", "XCLIPProcessor"),
+ ]
+ )
+
+PROCESSOR_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, PROCESSOR_MAPPING_NAMES)
+
+
+def processor_class_from_name(class_name: str):
+ for module_name, processors in PROCESSOR_MAPPING_NAMES.items():
+ if class_name in processors:
+ module_name = model_type_to_module_name(module_name)
+
+ module = importlib.import_module(f".{module_name}", "transformers.models")
+ try:
+ return getattr(module, class_name)
+ except AttributeError:
+ continue
+
+ for processor in PROCESSOR_MAPPING._extra_content.values():
+ if getattr(processor, "__name__", None) == class_name:
+ return processor
+
+ # We did not find the class, but maybe it's because a dep is missing. In that case, the class will be in the main
+ # init and we return the proper dummy to get an appropriate error message.
+ main_module = importlib.import_module("transformers")
+ if hasattr(main_module, class_name):
+ return getattr(main_module, class_name)
+
+ return None
+
+
+class AutoProcessor:
+ r"""
+ This is a generic processor class that will be instantiated as one of the processor classes of the library when
+ created with the [`AutoProcessor.from_pretrained`] class method.
+
+ This class cannot be instantiated directly using `__init__()` (throws an error).
+ """
+
+ def __init__(self):
+ raise OSError(
+ "AutoProcessor is designed to be instantiated "
+ "using the `AutoProcessor.from_pretrained(pretrained_model_name_or_path)` method."
+ )
+
+ @classmethod
+ @replace_list_option_in_docstrings(PROCESSOR_MAPPING_NAMES)
+ def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+ r"""
+ Instantiate one of the processor classes of the library from a pretrained model vocabulary.
+
+ The processor class to instantiate is selected based on the `model_type` property of the config object (either
+ passed as an argument or loaded from `pretrained_model_name_or_path` if possible):
+
+ List options
+
+ Params:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained feature_extractor hosted inside a model repo on
+ huggingface.co.
+ - a path to a *directory* containing a processor files saved using the `save_pretrained()` method,
+ e.g., `./my_model_directory/`.
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model feature extractor should be cached if the
+ standard cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the feature extractor files and override the cached versions
+ if they exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `hf auth login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ return_unused_kwargs (`bool`, *optional*, defaults to `False`):
+ If `False`, then this function returns just the final feature extractor object. If `True`, then this
+ functions returns a `Tuple(feature_extractor, unused_kwargs)` where *unused_kwargs* is a dictionary
+ consisting of the key/value pairs whose keys are not feature extractor attributes: i.e., the part of
+ `kwargs` which has not been used to update `feature_extractor` and is otherwise ignored.
+ trust_remote_code (`bool`, *optional*, defaults to `False`):
+ Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
+ should only be set to `True` for repositories you trust and in which you have read the code, as it will
+ execute code present on the Hub on your local machine.
+ kwargs (`dict[str, Any]`, *optional*):
+ The values in kwargs of any keys which are feature extractor attributes will be used to override the
+ loaded values. Behavior concerning key/value pairs whose keys are *not* feature extractor attributes is
+ controlled by the `return_unused_kwargs` keyword parameter.
+
+
+
+ Passing `token=True` is required when you want to use a private model.
+
+
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoProcessor
+
+ >>> # Download processor from huggingface.co and cache.
+ >>> processor = AutoProcessor.from_pretrained("facebook/wav2vec2-base-960h")
+
+ >>> # If processor files are in a directory (e.g. processor was saved using *save_pretrained('./test/saved_model/')*)
+ >>> # processor = AutoProcessor.from_pretrained("./test/saved_model/")
+ ```"""
+ config = kwargs.pop("config", None)
+ trust_remote_code = kwargs.pop("trust_remote_code", None)
+ kwargs["_from_auto"] = True
+
+ processor_class = None
+ processor_auto_map = None
+
+ # First, let's see if we have a processor or preprocessor config.
+ # Filter the kwargs for `cached_file`.
+ _hub_valid_kwargs = (
+ "cache_dir",
+ "force_download",
+ "proxies",
+ "token",
+ "revision",
+ "local_files_only",
+ "subfolder",
+ "repo_type",
+ "user_agent",
+ )
+ cached_file_kwargs = {key: kwargs[key] for key in _hub_valid_kwargs if key in kwargs}
+ # We don't want to raise
+ cached_file_kwargs.update(
+ {
+ "_raise_exceptions_for_gated_repo": False,
+ "_raise_exceptions_for_missing_entries": False,
+ "_raise_exceptions_for_connection_errors": False,
+ }
+ )
+
+ # Let's start by checking whether the processor class is saved in a processor config
+ processor_config_file = cached_file(pretrained_model_name_or_path, PROCESSOR_NAME, **cached_file_kwargs)
+ if processor_config_file is not None:
+ config_dict, _ = ProcessorMixin.get_processor_dict(pretrained_model_name_or_path, **kwargs)
+ processor_class = config_dict.get("processor_class")
+ if "AutoProcessor" in config_dict.get("auto_map", {}):
+ processor_auto_map = config_dict["auto_map"]["AutoProcessor"]
+
+ if processor_class is None:
+ # If not found, let's check whether the processor class is saved in an image processor config
+ preprocessor_config_file = cached_file(
+ pretrained_model_name_or_path, FEATURE_EXTRACTOR_NAME, **cached_file_kwargs
+ )
+ if preprocessor_config_file is not None:
+ config_dict, _ = ImageProcessingMixin.get_image_processor_dict(pretrained_model_name_or_path, **kwargs)
+ processor_class = config_dict.get("processor_class", None)
+ if "AutoProcessor" in config_dict.get("auto_map", {}):
+ processor_auto_map = config_dict["auto_map"]["AutoProcessor"]
+
+ # Saved as video processor
+ if preprocessor_config_file is None:
+ preprocessor_config_file = cached_file(
+ pretrained_model_name_or_path, VIDEO_PROCESSOR_NAME, **cached_file_kwargs
+ )
+ if preprocessor_config_file is not None:
+ config_dict, _ = BaseVideoProcessor.get_video_processor_dict(
+ pretrained_model_name_or_path, **kwargs
+ )
+ processor_class = config_dict.get("processor_class", None)
+ if "AutoProcessor" in config_dict.get("auto_map", {}):
+ processor_auto_map = config_dict["auto_map"]["AutoProcessor"]
+ # Saved as feature extractor
+ if preprocessor_config_file is None:
+ preprocessor_config_file = cached_file(
+ pretrained_model_name_or_path, FEATURE_EXTRACTOR_NAME, **cached_file_kwargs
+ )
+ if preprocessor_config_file is not None and processor_class is None:
+ config_dict, _ = FeatureExtractionMixin.get_feature_extractor_dict(
+ pretrained_model_name_or_path, **kwargs
+ )
+ processor_class = config_dict.get("processor_class", None)
+ if "AutoProcessor" in config_dict.get("auto_map", {}):
+ processor_auto_map = config_dict["auto_map"]["AutoProcessor"]
+
+ if processor_class is None:
+ # Next, let's check whether the processor class is saved in a tokenizer
+ tokenizer_config_file = cached_file(
+ pretrained_model_name_or_path, TOKENIZER_CONFIG_FILE, **cached_file_kwargs
+ )
+ if tokenizer_config_file is not None:
+ with open(tokenizer_config_file, encoding="utf-8") as reader:
+ config_dict = json.load(reader)
+
+ processor_class = config_dict.get("processor_class", None)
+ if "AutoProcessor" in config_dict.get("auto_map", {}):
+ processor_auto_map = config_dict["auto_map"]["AutoProcessor"]
+
+ if processor_class is None:
+ # Last resort: try loading the model config to get processor_class.
+ # This handles cases where processor info is only in config.json (not in any
+ # preprocessor/tokenizer config files). AutoConfig.from_pretrained may raise
+ # ValueError if the model_type is unrecognized or the config is invalid -
+ # we catch and ignore this to allow fallback to AutoTokenizer/AutoImageProcessor.
+ try:
+ if not isinstance(config, PreTrainedConfig):
+ config = AutoConfig.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+
+ processor_class = getattr(config, "processor_class", None)
+ if hasattr(config, "auto_map") and "AutoProcessor" in config.auto_map:
+ processor_auto_map = config.auto_map["AutoProcessor"]
+ except ValueError:
+ # Config loading failed (unrecognized model_type, invalid config, etc.)
+ # Continue to fallback logic below (AutoTokenizer, AutoImageProcessor, etc.)
+ pass
+
+ if processor_class is not None:
+ processor_class = processor_class_from_name(processor_class)
+
+ has_remote_code = processor_auto_map is not None
+ has_local_code = processor_class is not None or type(config) in PROCESSOR_MAPPING
+ explicit_local_code = has_local_code and not (
+ processor_class or PROCESSOR_MAPPING[type(config)]
+ ).__module__.startswith("transformers.")
+ if has_remote_code:
+ if "--" in processor_auto_map:
+ upstream_repo = processor_auto_map.split("--")[0]
+ else:
+ upstream_repo = None
+ trust_remote_code = resolve_trust_remote_code(
+ trust_remote_code, pretrained_model_name_or_path, has_local_code, has_remote_code, upstream_repo
+ )
+
+ if has_remote_code and trust_remote_code and not explicit_local_code:
+ processor_class = get_class_from_dynamic_module(
+ processor_auto_map, pretrained_model_name_or_path, **kwargs
+ )
+ _ = kwargs.pop("code_revision", None)
+ processor_class.register_for_auto_class()
+ return processor_class.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+ elif processor_class is not None:
+ return processor_class.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+ # Last try: we use the PROCESSOR_MAPPING.
+ elif type(config) in PROCESSOR_MAPPING:
+ return PROCESSOR_MAPPING[type(config)].from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+ # At this stage, there doesn't seem to be a `Processor` class available for this model.
+ # Let's try the commonly available classes
+ for klass in (AutoTokenizer, AutoImageProcessor, AutoVideoProcessor, AutoFeatureExtractor):
+ try:
+ return klass.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+ except Exception:
+ continue
+
+ raise ValueError(
+ f"Unrecognized processing class in {pretrained_model_name_or_path}. Can't instantiate a processor, a "
+ "tokenizer, an image processor, a video processor or a feature extractor for this model. "
+ "Make sure the repository contains the files of at least one of those processing classes."
+ )
+
+ @staticmethod
+ def register(config_class, processor_class, exist_ok=False):
+ """
+ Register a new processor for this class.
+
+ Args:
+ config_class ([`PreTrainedConfig`]):
+ The configuration corresponding to the model to register.
+ processor_class ([`ProcessorMixin`]): The processor to register.
+ """
+ PROCESSOR_MAPPING.register(config_class, processor_class, exist_ok=exist_ok)
+
+
+__all__ = ["PROCESSOR_MAPPING", "AutoProcessor"]
diff --git a/third_party/transformers/src/transformers/models/auto/tokenization_auto.py b/third_party/transformers/src/transformers/models/auto/tokenization_auto.py
new file mode 100644
index 0000000000000000000000000000000000000000..1b38f2e7a3f1f5adcb072025c4ae6dd783968238
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/auto/tokenization_auto.py
@@ -0,0 +1,871 @@
+# Copyright 2018 The HuggingFace Inc. team.
+#
+# 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.
+"""Auto Tokenizer class."""
+
+import importlib
+import json
+import os
+import sys
+from collections import OrderedDict
+from typing import Any
+
+from transformers.utils.import_utils import is_mistral_common_available
+
+from ...configuration_utils import PreTrainedConfig
+from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
+from ...modeling_gguf_pytorch_utils import load_gguf_checkpoint
+from ...tokenization_utils_base import TOKENIZER_CONFIG_FILE
+from ...utils import (
+ extract_commit_hash,
+ is_g2p_en_available,
+ is_sentencepiece_available,
+ is_tokenizers_available,
+ logging,
+)
+from ...utils.hub import cached_file
+from ..encoder_decoder import EncoderDecoderConfig
+from .auto_factory import _LazyAutoMapping
+from .configuration_auto import (
+ CONFIG_MAPPING_NAMES,
+ AutoConfig,
+ config_class_to_model_type,
+ model_type_to_module_name,
+ replace_list_option_in_docstrings,
+)
+
+
+if is_tokenizers_available():
+ from ...tokenization_utils_tokenizers import TokenizersBackend
+else:
+ TokenizersBackend = None
+
+if is_sentencepiece_available():
+ from ...tokenization_utils_sentencepiece import SentencePieceBackend
+else:
+ SentencePieceBackend = None
+
+logger = logging.get_logger(__name__)
+
+# V5: Simplified mapping - single tokenizer class per model type (always prefer tokenizers-based)
+REGISTERED_TOKENIZER_CLASSES: dict[str, type[Any]] = {}
+REGISTERED_FAST_ALIASES: dict[str, type[Any]] = {}
+
+TOKENIZER_MAPPING_NAMES = OrderedDict[str, str | None](
+ [
+ ("aimv2", "CLIPTokenizer" if is_tokenizers_available() else None),
+ ("albert", "AlbertTokenizer" if is_tokenizers_available() else None),
+ ("align", "BertTokenizer" if is_tokenizers_available() else None),
+ ("audioflamingo3", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("aya_vision", "CohereTokenizer" if is_tokenizers_available() else None),
+ ("bark", "BertTokenizer" if is_tokenizers_available() else None),
+ ("bart", "RobertaTokenizer" if is_tokenizers_available() else None),
+ ("barthez", "BarthezTokenizer" if is_tokenizers_available() else None),
+ ("bartpho", "BartphoTokenizer"),
+ ("bert", "BertTokenizer" if is_tokenizers_available() else None),
+ ("bert-generation", "BertGenerationTokenizer" if is_sentencepiece_available() else None),
+ ("bert-japanese", "BertJapaneseTokenizer"),
+ ("bertweet", "BertweetTokenizer"),
+ ("big_bird", "BigBirdTokenizer" if is_tokenizers_available() else None),
+ ("bigbird_pegasus", "PegasusTokenizer" if is_tokenizers_available() else None),
+ ("biogpt", "BioGptTokenizer"),
+ ("blenderbot", "BlenderbotTokenizer" if is_tokenizers_available() else None),
+ ("blenderbot-small", "BlenderbotSmallTokenizer"),
+ ("blip", "BertTokenizer" if is_tokenizers_available() else None),
+ ("blip-2", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("bridgetower", "RobertaTokenizer"),
+ ("bros", "BertTokenizer" if is_tokenizers_available() else None),
+ ("byt5", "ByT5Tokenizer"),
+ ("camembert", "CamembertTokenizer" if is_tokenizers_available() else None),
+ ("canine", "CanineTokenizer"),
+ ("chinese_clip", "BertTokenizer" if is_tokenizers_available() else None),
+ ("clap", "RobertaTokenizer"),
+ ("clip", "CLIPTokenizer" if is_tokenizers_available() else None),
+ ("clipseg", "CLIPTokenizer" if is_tokenizers_available() else None),
+ ("clvp", "ClvpTokenizer"),
+ ("code_llama", "CodeLlamaTokenizer" if is_tokenizers_available() else None),
+ ("codegen", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("cohere", "CohereTokenizer" if is_tokenizers_available() else None),
+ ("cohere2", "CohereTokenizer" if is_tokenizers_available() else None),
+ ("colqwen2", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("convbert", "BertTokenizer" if is_tokenizers_available() else None),
+ ("cpm", "CpmTokenizer" if is_tokenizers_available() else None),
+ ("cpmant", "CpmAntTokenizer"),
+ ("ctrl", "CTRLTokenizer"),
+ ("data2vec-audio", "Wav2Vec2CTCTokenizer"),
+ ("data2vec-text", "RobertaTokenizer"),
+ ("dbrx", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("deberta", "DebertaTokenizer" if is_tokenizers_available() else None),
+ ("deberta-v2", "DebertaV2Tokenizer" if is_tokenizers_available() else None),
+ ("dia", "DiaTokenizer"),
+ ("distilbert", "BertTokenizer" if is_tokenizers_available() else None),
+ ("dpr", "DPRQuestionEncoderTokenizer" if is_tokenizers_available() else None),
+ ("electra", "BertTokenizer" if is_tokenizers_available() else None),
+ ("emu3", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("ernie", "BertTokenizer" if is_tokenizers_available() else None),
+ ("esm", "EsmTokenizer"),
+ ("falcon_mamba", "GPTNeoXTokenizer" if is_tokenizers_available() else None),
+ ("fastspeech2_conformer", "FastSpeech2ConformerTokenizer" if is_g2p_en_available() else None),
+ ("flaubert", "FlaubertTokenizer"),
+ ("flava", "BertTokenizer" if is_tokenizers_available() else None),
+ ("flex_olmo", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("florence2", "BartTokenizer" if is_tokenizers_available() else None),
+ ("fnet", "FNetTokenizer" if is_tokenizers_available() else None),
+ ("fsmt", "FSMTTokenizer"),
+ ("funnel", "FunnelTokenizer" if is_tokenizers_available() else None),
+ ("gemma", "GemmaTokenizer" if is_tokenizers_available() else None),
+ ("gemma2", "GemmaTokenizer" if is_tokenizers_available() else None),
+ ("gemma3", "GemmaTokenizer" if is_tokenizers_available() else None),
+ ("gemma3_text", "GemmaTokenizer" if is_tokenizers_available() else None),
+ ("gemma3n", "GemmaTokenizer" if is_tokenizers_available() else None),
+ ("gemma3n_text", "GemmaTokenizer" if is_tokenizers_available() else None),
+ ("git", "BertTokenizer" if is_tokenizers_available() else None),
+ ("glm", "TokenizersBackend" if is_tokenizers_available() else None),
+ ("glm4", "TokenizersBackend" if is_tokenizers_available() else None),
+ ("glm4_moe", "TokenizersBackend" if is_tokenizers_available() else None),
+ ("glm4_moe_lite", "TokenizersBackend" if is_tokenizers_available() else None),
+ ("glm4v", "TokenizersBackend" if is_tokenizers_available() else None),
+ ("glm4v_moe", "TokenizersBackend" if is_tokenizers_available() else None),
+ ("glm_image", "TokenizersBackend" if is_tokenizers_available() else None),
+ ("glmasr", "TokenizersBackend" if is_tokenizers_available() else None),
+ ("got_ocr2", "TokenizersBackend" if is_tokenizers_available() else None),
+ ("gpt-sw3", "GPTSw3Tokenizer" if is_sentencepiece_available() else None),
+ ("gpt2", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("gpt_bigcode", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("gpt_neo", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("gpt_neox", "GPTNeoXTokenizer" if is_tokenizers_available() else None),
+ ("gpt_neox_japanese", "GPTNeoXJapaneseTokenizer"),
+ ("gptj", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("granite", "GPT2Tokenizer"),
+ ("granitemoe", "GPT2Tokenizer"),
+ ("granitemoehybrid", "GPT2Tokenizer"),
+ ("granitemoeshared", "GPT2Tokenizer"),
+ ("grounding-dino", "BertTokenizer" if is_tokenizers_available() else None),
+ ("groupvit", "CLIPTokenizer" if is_tokenizers_available() else None),
+ ("herbert", "HerbertTokenizer" if is_tokenizers_available() else None),
+ ("hubert", "Wav2Vec2CTCTokenizer"),
+ ("ibert", "RobertaTokenizer"),
+ ("idefics", "LlamaTokenizer" if is_tokenizers_available() else None),
+ ("idefics2", "LlamaTokenizer" if is_tokenizers_available() else None),
+ ("instructblip", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("instructblipvideo", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("internvl", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("jais2", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("jina_embeddings_v3", "XLMRobertaTokenizer" if is_tokenizers_available() else None),
+ ("kosmos-2", "XLMRobertaTokenizer" if is_tokenizers_available() else None),
+ ("lasr_ctc", "LasrTokenizer" if is_tokenizers_available() else None),
+ ("lasr_encoder", "LasrTokenizer" if is_tokenizers_available() else None),
+ ("layoutlm", "BertTokenizer" if is_tokenizers_available() else None),
+ ("layoutlmv2", "LayoutLMv2Tokenizer" if is_tokenizers_available() else None),
+ ("layoutlmv3", "LayoutLMv3Tokenizer" if is_tokenizers_available() else None),
+ ("layoutxlm", "LayoutXLMTokenizer" if is_tokenizers_available() else None),
+ ("led", "LEDTokenizer" if is_tokenizers_available() else None),
+ ("lighton_ocr", "Qwen2TokenizerFast" if is_tokenizers_available() else None),
+ ("lilt", "RobertaTokenizer" if is_tokenizers_available() else None),
+ ("longformer", "RobertaTokenizer" if is_tokenizers_available() else None),
+ ("luke", "LukeTokenizer"),
+ ("lxmert", "LxmertTokenizer" if is_tokenizers_available() else None),
+ ("m2m_100", "M2M100Tokenizer" if is_sentencepiece_available() else None),
+ ("mamba", "GPTNeoXTokenizer" if is_tokenizers_available() else None),
+ ("mamba2", "GPTNeoXTokenizer" if is_tokenizers_available() else None),
+ ("marian", "MarianTokenizer" if is_sentencepiece_available() else None),
+ ("markuplm", "MarkupLMTokenizer" if is_tokenizers_available() else None),
+ ("mbart", "MBartTokenizer" if is_tokenizers_available() else None),
+ ("mbart50", "MBart50Tokenizer" if is_tokenizers_available() else None),
+ ("mega", "RobertaTokenizer"),
+ ("megatron-bert", "BertTokenizer" if is_tokenizers_available() else None),
+ ("metaclip_2", "XLMRobertaTokenizer" if is_tokenizers_available() else None),
+ ("mgp-str", "MgpstrTokenizer"),
+ (
+ "ministral",
+ "MistralCommonBackend"
+ if is_mistral_common_available()
+ else ("TokenizersBackend" if is_tokenizers_available() else None),
+ ),
+ (
+ "ministral3",
+ "MistralCommonBackend"
+ if is_mistral_common_available()
+ else ("TokenizersBackend" if is_tokenizers_available() else None),
+ ),
+ (
+ "mistral",
+ "MistralCommonBackend"
+ if is_mistral_common_available()
+ else ("TokenizersBackend" if is_tokenizers_available() else None),
+ ),
+ (
+ "mistral3",
+ "MistralCommonBackend"
+ if is_mistral_common_available()
+ else ("TokenizersBackend" if is_tokenizers_available() else None),
+ ),
+ (
+ "mixtral",
+ "MistralCommonBackend"
+ if is_mistral_common_available()
+ else ("TokenizersBackend" if is_tokenizers_available() else None),
+ ),
+ ("mluke", "MLukeTokenizer" if is_sentencepiece_available() else None),
+ ("mm-grounding-dino", "BertTokenizer" if is_tokenizers_available() else None),
+ ("mobilebert", "MobileBertTokenizer" if is_tokenizers_available() else None),
+ ("mpnet", "MPNetTokenizer" if is_tokenizers_available() else None),
+ ("mpt", "GPTNeoXTokenizer" if is_tokenizers_available() else None),
+ ("mra", "RobertaTokenizer"),
+ ("mt5", "T5Tokenizer" if is_tokenizers_available() else None),
+ ("musicgen", "T5Tokenizer" if is_tokenizers_available() else None),
+ ("musicgen_melody", "T5Tokenizer" if is_tokenizers_available() else None),
+ ("mvp", "MvpTokenizer" if is_tokenizers_available() else None),
+ ("myt5", "MyT5Tokenizer"),
+ ("nezha", "BertTokenizer" if is_tokenizers_available() else None),
+ ("nllb", "NllbTokenizer" if is_tokenizers_available() else None),
+ ("nllb-moe", "NllbTokenizer" if is_tokenizers_available() else None),
+ ("nomic_bert", "BertTokenizer" if is_tokenizers_available() else None),
+ ("nougat", "NougatTokenizer" if is_tokenizers_available() else None),
+ ("nystromformer", "AlbertTokenizer" if is_tokenizers_available() else None),
+ ("olmo", "GPTNeoXTokenizer" if is_tokenizers_available() else None),
+ ("olmo2", "GPTNeoXTokenizer" if is_tokenizers_available() else None),
+ ("olmo3", "TokenizersBackend" if is_tokenizers_available() else None),
+ ("olmo_hybrid", "TokenizersBackend" if is_tokenizers_available() else None),
+ ("olmoe", "GPTNeoXTokenizer" if is_tokenizers_available() else None),
+ ("omdet-turbo", "CLIPTokenizer" if is_tokenizers_available() else None),
+ ("oneformer", "CLIPTokenizer" if is_tokenizers_available() else None),
+ ("openai-gpt", "OpenAIGPTTokenizer" if is_tokenizers_available() else None),
+ ("opt", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("ovis2", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("owlv2", "CLIPTokenizer" if is_tokenizers_available() else None),
+ ("owlvit", "CLIPTokenizer" if is_tokenizers_available() else None),
+ ("pegasus", "PegasusTokenizer" if is_tokenizers_available() else None),
+ ("pegasus_x", "PegasusTokenizer" if is_tokenizers_available() else None),
+ ("perceiver", "PerceiverTokenizer"),
+ ("phi", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("phobert", "PhobertTokenizer"),
+ ("pix2struct", "T5Tokenizer" if is_tokenizers_available() else None),
+ (
+ "pixtral",
+ "MistralCommonBackend"
+ if is_mistral_common_available()
+ else ("TokenizersBackend" if is_tokenizers_available() else None),
+ ),
+ ("plbart", "PLBartTokenizer" if is_tokenizers_available() else None),
+ ("prophetnet", "ProphetNetTokenizer"),
+ ("qdqbert", "BertTokenizer" if is_tokenizers_available() else None),
+ ("qwen2", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("qwen2_5_omni", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("qwen2_5_vl", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("qwen2_audio", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("qwen2_moe", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("qwen2_vl", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("qwen3", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("qwen3_5", "Qwen3_5Tokenizer" if is_tokenizers_available() else None),
+ ("qwen3_5_moe", "Qwen3_5Tokenizer" if is_tokenizers_available() else None),
+ ("qwen3_moe", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("qwen3_next", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("qwen3_omni_moe", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("qwen3_vl", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("qwen3_vl_moe", "Qwen2Tokenizer" if is_tokenizers_available() else None),
+ ("rag", "RagTokenizer"),
+ ("realm", "BertTokenizer" if is_tokenizers_available() else None),
+ ("recurrent_gemma", "GemmaTokenizer" if is_tokenizers_available() else None),
+ ("reformer", "ReformerTokenizer" if is_tokenizers_available() else None),
+ ("rembert", "RemBertTokenizer" if is_tokenizers_available() else None),
+ ("retribert", "BertTokenizer" if is_tokenizers_available() else None),
+ ("roberta", "RobertaTokenizer"),
+ ("roberta-prelayernorm", "RobertaTokenizer"),
+ ("roc_bert", "RoCBertTokenizer"),
+ ("roformer", "RoFormerTokenizer" if is_tokenizers_available() else None),
+ ("rwkv", "GPTNeoXTokenizer" if is_tokenizers_available() else None),
+ ("sam3", "CLIPTokenizer" if is_tokenizers_available() else None),
+ ("sam3_video", "CLIPTokenizer" if is_tokenizers_available() else None),
+ ("seamless_m4t", "SeamlessM4TTokenizer" if is_tokenizers_available() else None),
+ ("seamless_m4t_v2", "SeamlessM4TTokenizer" if is_tokenizers_available() else None),
+ ("shieldgemma2", "GemmaTokenizer" if is_tokenizers_available() else None),
+ ("siglip", "SiglipTokenizer" if is_sentencepiece_available() else None),
+ ("siglip2", "Siglip2Tokenizer" if is_tokenizers_available() else None),
+ ("speech_to_text", "Speech2TextTokenizer" if is_sentencepiece_available() else None),
+ ("speecht5", "SpeechT5Tokenizer" if is_sentencepiece_available() else None),
+ ("splinter", "SplinterTokenizer"),
+ ("squeezebert", "BertTokenizer" if is_tokenizers_available() else None),
+ ("stablelm", "GPTNeoXTokenizer" if is_tokenizers_available() else None),
+ ("starcoder2", "GPT2Tokenizer" if is_tokenizers_available() else None),
+ ("switch_transformers", "T5Tokenizer" if is_tokenizers_available() else None),
+ ("t5", "T5Tokenizer" if is_tokenizers_available() else None),
+ ("t5gemma", "GemmaTokenizer" if is_tokenizers_available() else None),
+ ("tapas", "TapasTokenizer"),
+ ("trocr", "XLMRobertaTokenizer" if is_tokenizers_available() else None),
+ ("tvp", "BertTokenizer" if is_tokenizers_available() else None),
+ ("udop", "UdopTokenizer" if is_tokenizers_available() else None),
+ ("umt5", "T5Tokenizer" if is_tokenizers_available() else None),
+ ("unispeech", "Wav2Vec2CTCTokenizer"),
+ ("unispeech-sat", "Wav2Vec2CTCTokenizer"),
+ ("vilt", "BertTokenizer" if is_tokenizers_available() else None),
+ ("visual_bert", "BertTokenizer" if is_tokenizers_available() else None),
+ ("vits", "VitsTokenizer"),
+ (
+ "voxtral",
+ "MistralCommonBackend"
+ if is_mistral_common_available()
+ else ("TokenizersBackend" if is_tokenizers_available() else None),
+ ),
+ (
+ "voxtral_realtime",
+ "MistralCommonBackend"
+ if is_mistral_common_available()
+ else ("TokenizersBackend" if is_tokenizers_available() else None),
+ ),
+ ("wav2vec2", "Wav2Vec2CTCTokenizer"),
+ ("wav2vec2-bert", "Wav2Vec2CTCTokenizer"),
+ ("wav2vec2-conformer", "Wav2Vec2CTCTokenizer"),
+ ("wav2vec2_phoneme", "Wav2Vec2PhonemeCTCTokenizer"),
+ ("whisper", "WhisperTokenizer" if is_tokenizers_available() else None),
+ ("xclip", "CLIPTokenizer" if is_tokenizers_available() else None),
+ ("xglm", "XGLMTokenizer" if is_tokenizers_available() else None),
+ ("xlm", "XLMTokenizer"),
+ ("xlm-roberta", "XLMRobertaTokenizer" if is_tokenizers_available() else None),
+ ("xlm-roberta-xl", "XLMRobertaTokenizer" if is_tokenizers_available() else None),
+ ("xlnet", "XLNetTokenizer" if is_tokenizers_available() else None),
+ ("xlstm", "GPTNeoXTokenizer" if is_tokenizers_available() else None),
+ ("xmod", "XLMRobertaTokenizer" if is_tokenizers_available() else None),
+ ("yoso", "AlbertTokenizer" if is_tokenizers_available() else None),
+ ]
+)
+
+# Models with incorrect tokenizer_class in their Hub tokenizer_config.json files.
+# These models will be forced to use TokenizersBackend.
+MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS: set[str] = {
+ "arctic",
+ "chameleon",
+ "chatlm",
+ "deepseek_v2",
+ "deepseek_v3",
+ "deepseek_vl",
+ "deepseek_vl_hybrid",
+ "deepseek_vl_v2",
+ "fuyu",
+ "h2ovl_chat",
+ "hyperclovax_vlm",
+ "internlm2",
+ "internvl_chat",
+ "jamba",
+ "janus",
+ "kimi_k25",
+ "llava",
+ "llava_next",
+ "minicpmv",
+ "minimax_m2",
+ "modernbert",
+ "molmo",
+ "molmo2",
+ "nemotron",
+ "nvfp4",
+ "opencua",
+ "openvla",
+ "phi3",
+ "phi3_v",
+ "phimoe",
+ "step3p5",
+ "vipllava",
+ "cohere_asr",
+}
+
+for model_type in MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS:
+ if model_type not in TOKENIZER_MAPPING_NAMES:
+ TOKENIZER_MAPPING_NAMES[model_type] = "TokenizersBackend" if is_tokenizers_available() else None
+
+TOKENIZER_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, TOKENIZER_MAPPING_NAMES)
+
+CONFIG_TO_TYPE = {v: k for k, v in CONFIG_MAPPING_NAMES.items()}
+
+
+def load_vocab(vocab_file):
+ """Loads a vocabulary file into a dictionary."""
+ with open(vocab_file, "r", encoding="utf-8") as reader:
+ return json.load(reader)
+
+
+def load_merges(merges_file):
+ """Loads a merges file into a list."""
+ merges = []
+ with open(merges_file, "r", encoding="utf-8") as reader:
+ for line in reader:
+ line = line.strip()
+ if line and not line.startswith("#"):
+ merges.append(tuple(line.split()))
+ return merges
+
+
+def tokenizer_class_from_name(class_name: str) -> type[Any] | None:
+ # Bloom tokenizer classes were removed but should map to the fast backend for BC
+ if class_name in {"BloomTokenizer", "BloomTokenizerFast"}:
+ return TokenizersBackend
+
+ if class_name in REGISTERED_FAST_ALIASES:
+ return REGISTERED_FAST_ALIASES[class_name]
+
+ if class_name in REGISTERED_TOKENIZER_CLASSES:
+ return REGISTERED_TOKENIZER_CLASSES[class_name]
+
+ if class_name == "TokenizersBackend":
+ return TokenizersBackend
+
+ # V5: TOKENIZER_MAPPING_NAMES now maps to single strings, not tuples
+ for module_name, tokenizer_class in TOKENIZER_MAPPING_NAMES.items():
+ if tokenizer_class == class_name:
+ module_name = model_type_to_module_name(module_name)
+ if (
+ module_name in ["mistral", "mistral3", "mixtral", "ministral", "ministral3", "pixtral", "voxtral"]
+ and class_name == "MistralCommonBackend"
+ ):
+ module = importlib.import_module(".tokenization_mistral_common", "transformers")
+ else:
+ module = importlib.import_module(f".{module_name}", "transformers.models")
+ try:
+ result = getattr(module, class_name)
+ # BC v5: expose XxxFast alias and tokenization_*_fast submodule for pre-v5 remote code.
+ if (submod := getattr(result, "__module__", None)) and submod in sys.modules:
+ base_mod = sys.modules[submod]
+ setattr(base_mod, result.__name__ + "Fast", result)
+ sys.modules.setdefault(submod + "_fast", base_mod)
+ return result
+ except AttributeError:
+ continue
+
+ for tokenizer in TOKENIZER_MAPPING._extra_content.values():
+ if getattr(tokenizer, "__name__", None) == class_name:
+ return tokenizer
+
+ # We did not find the class, but maybe it's because a dep is missing. In that case, the class will be in the main
+ # We did not find the class, but maybe it's because a dep is missing. In that case, the class will be in the main
+ # init and we return the proper dummy to get an appropriate error message.
+ main_module = importlib.import_module("transformers")
+ if hasattr(main_module, class_name):
+ return getattr(main_module, class_name)
+
+ # BC v5: If a XxxFast class is not found, retry without 'Fast' for tokenizers saved pre-v5.
+ if class_name.endswith("Fast"):
+ return tokenizer_class_from_name(class_name[:-4])
+
+ return None
+
+
+def get_tokenizer_config(
+ pretrained_model_name_or_path: str | os.PathLike[str],
+ cache_dir: str | os.PathLike[str] | None = None,
+ force_download: bool = False,
+ proxies: dict[str, str] | None = None,
+ token: bool | str | None = None,
+ revision: str | None = None,
+ local_files_only: bool = False,
+ subfolder: str = "",
+ **kwargs,
+) -> dict[str, Any]:
+ """
+ Loads the tokenizer configuration from a pretrained model tokenizer configuration.
+
+ Args:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained model configuration hosted inside a model repo on
+ huggingface.co.
+ - a path to a *directory* containing a configuration file saved using the
+ [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
+
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
+ cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the configuration files and override the cached versions if they
+ exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `hf auth login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ local_files_only (`bool`, *optional*, defaults to `False`):
+ If `True`, will only try to load the tokenizer configuration from local files.
+ subfolder (`str`, *optional*, defaults to `""`):
+ In case the tokenizer config is located inside a subfolder of the model repo on huggingface.co, you can
+ specify the folder name here.
+
+
+
+ Passing `token=True` is required when you want to use a private model.
+
+
+
+ Returns:
+ `dict`: The configuration of the tokenizer.
+
+ Examples:
+
+ ```python
+ # Download configuration from huggingface.co and cache.
+ tokenizer_config = get_tokenizer_config("google-bert/bert-base-uncased")
+ # This model does not have a tokenizer config so the result will be an empty dict.
+ tokenizer_config = get_tokenizer_config("FacebookAI/xlm-roberta-base")
+
+ # Save a pretrained tokenizer locally and you can reload its config
+ from transformers import AutoTokenizer
+
+ tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-cased")
+ tokenizer.save_pretrained("tokenizer-test")
+ tokenizer_config = get_tokenizer_config("tokenizer-test")
+ ```"""
+ commit_hash = kwargs.get("_commit_hash")
+ resolved_config_file = cached_file(
+ pretrained_model_name_or_path,
+ TOKENIZER_CONFIG_FILE,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ token=token,
+ revision=revision,
+ local_files_only=local_files_only,
+ subfolder=subfolder,
+ _raise_exceptions_for_gated_repo=False,
+ _raise_exceptions_for_missing_entries=False,
+ _raise_exceptions_for_connection_errors=False,
+ _commit_hash=commit_hash,
+ )
+ if resolved_config_file is None:
+ logger.info("Could not locate the tokenizer configuration file, will try to use the model config instead.")
+ return {}
+ commit_hash = extract_commit_hash(resolved_config_file, commit_hash)
+
+ with open(resolved_config_file, encoding="utf-8") as reader:
+ result = json.load(reader)
+ result["_commit_hash"] = commit_hash
+ return result
+
+
+class AutoTokenizer:
+ r"""
+ This is a generic tokenizer class that will be instantiated as one of the tokenizer classes of the library when
+ created with the [`AutoTokenizer.from_pretrained`] class method.
+
+ This class cannot be instantiated directly using `__init__()` (throws an error).
+ """
+
+ def __init__(self):
+ raise OSError(
+ "AutoTokenizer is designed to be instantiated "
+ "using the `AutoTokenizer.from_pretrained(pretrained_model_name_or_path)` method."
+ )
+
+ @classmethod
+ @replace_list_option_in_docstrings(TOKENIZER_MAPPING_NAMES)
+ def from_pretrained(
+ cls, pretrained_model_name_or_path, *inputs, **kwargs
+ ) -> TokenizersBackend | SentencePieceBackend:
+ r"""
+ Instantiate one of the tokenizer classes of the library from a pretrained model vocabulary.
+
+ The tokenizer class to instantiate is selected based on the `model_type` property of the config object (either
+ passed as an argument or loaded from `pretrained_model_name_or_path` if possible), or when it's missing, by
+ falling back to using pattern matching on `pretrained_model_name_or_path`:
+
+ List options
+
+ Params:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ Can be either:
+
+ - A string, the *model id* of a predefined tokenizer hosted inside a model repo on huggingface.co.
+ - A path to a *directory* containing vocabulary files required by the tokenizer, for instance saved
+ using the [`~PreTrainedTokenizer.save_pretrained`] method, e.g., `./my_model_directory/`.
+ - a path to a single saved vocabulary file if and only if the tokenizer only requires a
+ single vocabulary file (like Bert or XLNet), e.g.: `./my_model_directory/vocab.txt`. (Not
+ applicable to all derived classes)
+ inputs (additional positional arguments, *optional*):
+ Will be passed along to the Tokenizer `__init__()` method.
+ config ([`PreTrainedConfig`], *optional*)
+ The configuration object used to determine the tokenizer class to instantiate.
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the
+ standard cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force the (re-)download the model weights and configuration files and override the
+ cached versions if they exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ subfolder (`str`, *optional*):
+ In case the relevant files are located inside a subfolder of the model repo on huggingface.co (e.g. for
+ facebook/rag-token-base), specify it here.
+ tokenizer_type (`str`, *optional*):
+ Tokenizer type to be loaded.
+ backend (`str`, *optional*, defaults to `"tokenizers"`):
+ Backend to use for tokenization. Valid options are:
+ - `"tokenizers"`: Use the HuggingFace tokenizers library backend (default)
+ - `"sentencepiece"`: Use the SentencePiece backend
+ trust_remote_code (`bool`, *optional*, defaults to `False`):
+ Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
+ should only be set to `True` for repositories you trust and in which you have read the code, as it will
+ execute code present on the Hub on your local machine.
+ kwargs (additional keyword arguments, *optional*):
+ Will be passed to the Tokenizer `__init__()` method. Can be used to set special tokens like
+ `bos_token`, `eos_token`, `unk_token`, `sep_token`, `pad_token`, `cls_token`, `mask_token`,
+ `additional_special_tokens`. See parameters in the `__init__()` for more details.
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer
+
+ >>> # Download vocabulary from huggingface.co and cache.
+ >>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
+
+ >>> # Download vocabulary from huggingface.co (user-uploaded) and cache.
+ >>> tokenizer = AutoTokenizer.from_pretrained("dbmdz/bert-base-german-cased")
+
+ >>> # If vocabulary files are in a directory (e.g. tokenizer was saved using *save_pretrained('./test/saved_model/')*)
+ >>> # tokenizer = AutoTokenizer.from_pretrained("./test/bert_saved_model/")
+
+ >>> # Download vocabulary from huggingface.co and define model-specific arguments
+ >>> tokenizer = AutoTokenizer.from_pretrained("FacebookAI/roberta-base", add_prefix_space=True)
+
+ >>> # Explicitly use the tokenizers backend
+ >>> tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer", backend="tokenizers")
+
+ >>> # Explicitly use the sentencepiece backend
+ >>> tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer", backend="sentencepiece")
+ ```"""
+ config = kwargs.pop("config", None)
+ kwargs["_from_auto"] = True
+
+ # V5: Always use fast tokenizers, ignore use_fast parameter
+ _ = kwargs.pop("use_fast", None)
+ tokenizer_type = kwargs.pop("tokenizer_type", None)
+ trust_remote_code = kwargs.pop("trust_remote_code", None)
+ gguf_file = kwargs.get("gguf_file")
+
+ # First, let's see whether the tokenizer_type is passed so that we can leverage it
+ if tokenizer_type is not None:
+ tokenizer_class_name = TOKENIZER_MAPPING_NAMES.get(tokenizer_type, None)
+
+ if tokenizer_class_name is None:
+ raise ValueError(
+ f"Passed `tokenizer_type` {tokenizer_type} does not exist. `tokenizer_type` should be one of "
+ f"{', '.join(c for c in TOKENIZER_MAPPING_NAMES)}."
+ )
+
+ tokenizer_class = tokenizer_class_from_name(tokenizer_class_name)
+
+ if tokenizer_class is None:
+ raise ValueError(f"Tokenizer class {tokenizer_class_name} is not currently imported.")
+
+ return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+
+ if gguf_file:
+ gguf_path = cached_file(pretrained_model_name_or_path, gguf_file, **kwargs)
+ config_dict = load_gguf_checkpoint(gguf_path, return_tensors=False)["config"]
+ config = AutoConfig.for_model(**config_dict)
+ elif config is None:
+ try:
+ config = AutoConfig.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+ except (ValueError, OSError):
+ config = PreTrainedConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)
+
+ config_model_type = config.model_type
+
+ # Next, let's try to use the tokenizer_config file to get the tokenizer class.
+ tokenizer_config = get_tokenizer_config(pretrained_model_name_or_path, **kwargs)
+ tokenizer_config_class = tokenizer_config.get("tokenizer_class", None)
+
+ # Check for auto_map early to handle dynamic tokenizers properly
+ tokenizer_auto_map = None
+ if "auto_map" in tokenizer_config:
+ if isinstance(tokenizer_config["auto_map"], (tuple, list)):
+ # Legacy format for dynamic tokenizers
+ tokenizer_auto_map = tokenizer_config["auto_map"]
+ else:
+ tokenizer_auto_map = tokenizer_config["auto_map"].get("AutoTokenizer", None)
+
+ # if there is a config, we can check that the tokenizer class != than model class and can thus assume we need to use TokenizersBackend
+ # Skip this early exit if auto_map is present (custom tokenizer with trust_remote_code)
+ if (
+ tokenizer_auto_map is None
+ and tokenizer_config_class is not None
+ and config_model_type is not None
+ and config_model_type != ""
+ and TOKENIZER_MAPPING_NAMES.get(config_model_type) is not None
+ and (TOKENIZER_MAPPING_NAMES.get(config_model_type).removesuffix("Fast"))
+ != (tokenizer_config_class.removesuffix("Fast"))
+ ):
+ # new model, but we ignore it unless the model type is the same
+ if TokenizersBackend is not None:
+ try:
+ return TokenizersBackend.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+ except Exception as e:
+ logger.debug(f"Failed to use TokenizersBackend: {e}")
+
+ return tokenizer_class_from_name(tokenizer_config_class).from_pretrained(
+ pretrained_model_name_or_path, *inputs, **kwargs
+ )
+
+ if "_commit_hash" in tokenizer_config:
+ kwargs["_commit_hash"] = tokenizer_config["_commit_hash"]
+
+ if tokenizer_config_class and tokenizer_config_class.endswith("Fast"):
+ tokenizer_config_class = tokenizer_config_class[:-4]
+
+ has_remote_code = tokenizer_auto_map is not None
+ has_local_code = type(config) in TOKENIZER_MAPPING or (
+ tokenizer_config_class is not None
+ and (
+ tokenizer_class_from_name(tokenizer_config_class) is not None
+ or tokenizer_class_from_name(tokenizer_config_class + "Fast") is not None
+ )
+ )
+ explicit_local_code = (
+ has_local_code
+ and type(config) not in TOKENIZER_MAPPING
+ and (
+ tokenizer_config_class is not None
+ and not (
+ tokenizer_class_from_name(tokenizer_config_class)
+ or tokenizer_class_from_name(tokenizer_config_class + "Fast")
+ ).__module__.startswith("transformers.")
+ )
+ )
+ # V5: Skip remote tokenizer for custom models with incorrect hub tokenizer class
+ if has_remote_code and config_model_type in MODELS_WITH_INCORRECT_HUB_TOKENIZER_CLASS:
+ has_remote_code = False
+ tokenizer_auto_map = None
+
+ if has_remote_code:
+ # V5: Always prefer fast tokenizer (index 1), fallback to slow (index 0)
+ if tokenizer_auto_map[1] is not None:
+ class_ref = tokenizer_auto_map[1]
+ else:
+ class_ref = tokenizer_auto_map[0]
+ if "--" in class_ref:
+ upstream_repo = class_ref.split("--")[0]
+ else:
+ upstream_repo = None
+ trust_remote_code = resolve_trust_remote_code(
+ trust_remote_code, pretrained_model_name_or_path, has_local_code, has_remote_code, upstream_repo
+ )
+
+ if has_remote_code and trust_remote_code and not explicit_local_code:
+ # BC v5: register *Fast aliases before remote code loads.
+ if tokenizer_config_class:
+ tokenizer_class_from_name(tokenizer_config_class.removesuffix("Fast"))
+ tokenizer_class = get_class_from_dynamic_module(class_ref, pretrained_model_name_or_path, **kwargs)
+ _ = kwargs.pop("code_revision", None)
+ tokenizer_class.register_for_auto_class()
+ return tokenizer_class.from_pretrained(
+ pretrained_model_name_or_path, *inputs, trust_remote_code=trust_remote_code, **kwargs
+ )
+ elif tokenizer_config_class is not None:
+ tokenizer_class_candidate = tokenizer_config_class
+ tokenizer_class = tokenizer_class_from_name(tokenizer_class_candidate)
+ if tokenizer_class is None and not tokenizer_class_candidate.endswith("Fast"):
+ tokenizer_class = tokenizer_class_from_name(tokenizer_class_candidate + "Fast")
+ if tokenizer_class is not None and tokenizer_class.__name__ == "PythonBackend":
+ tokenizer_class = TokenizersBackend
+ # Fallback to TokenizersBackend if the class wasn't found
+ if tokenizer_class is None:
+ tokenizer_class = TokenizersBackend
+
+ return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+ elif getattr(config, "tokenizer_class", None):
+ _class = config.tokenizer_class
+ if "PreTrainedTokenizerFast" not in _class and _class.endswith("Fast"):
+ _class = _class[:-4]
+ tokenizer_class = tokenizer_class_from_name(_class)
+ return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+
+ # Otherwise we have to be creative.
+ # if model is an encoder decoder, the encoder tokenizer class is used by default
+ if isinstance(config, EncoderDecoderConfig):
+ if type(config.decoder) is not type(config.encoder):
+ logger.warning(
+ f"The encoder model config class: {config.encoder.__class__} is different from the decoder model "
+ f"config class: {config.decoder.__class__}. It is not recommended to use the "
+ "`AutoTokenizer.from_pretrained()` method in this case. Please use the encoder and decoder "
+ "specific tokenizer classes."
+ )
+ config = config.encoder
+
+ model_type = config_class_to_model_type(type(config).__name__) or getattr(config, "model_type", None)
+ if model_type is not None:
+ tokenizer_class = TOKENIZER_MAPPING.get(type(config), TokenizersBackend)
+ if tokenizer_class is not None:
+ return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+
+ # Fallback: try tokenizer_class from tokenizer_config.json
+ tokenizer_config_class = tokenizer_config.get("tokenizer_class", None)
+ if tokenizer_config_class is not None:
+ if tokenizer_config_class != "TokenizersBackend" and tokenizer_config_class.endswith("Fast"):
+ tokenizer_config_class = tokenizer_config_class[:-4]
+ tokenizer_class = tokenizer_class_from_name(tokenizer_config_class)
+ if tokenizer_class is None and not tokenizer_config_class.endswith("Fast"):
+ tokenizer_class = tokenizer_class_from_name(tokenizer_config_class + "Fast")
+ if tokenizer_class is not None and tokenizer_class.__name__ == "PythonBackend":
+ tokenizer_class = TokenizersBackend
+ if tokenizer_class is None:
+ tokenizer_class = TokenizersBackend
+ return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+
+ raise ValueError(
+ f"Unrecognized configuration class {config.__class__} to build an AutoTokenizer.\n"
+ f"Model type should be one of {', '.join(c.__name__ for c in TOKENIZER_MAPPING)}."
+ )
+
+ @staticmethod
+ def register(
+ config_class, tokenizer_class=None, slow_tokenizer_class=None, fast_tokenizer_class=None, exist_ok=False
+ ):
+ """
+ Register a new tokenizer in this mapping.
+
+ Args:
+ config_class ([`PreTrainedConfig`]):
+ The configuration corresponding to the model to register.
+ tokenizer_class: The tokenizer class to register (V5 - preferred parameter).
+ slow_tokenizer_class: (Deprecated) The slow tokenizer to register.
+ fast_tokenizer_class: (Deprecated) The fast tokenizer to register.
+ """
+ if tokenizer_class is None:
+ # Legacy: prefer fast over slow
+ if fast_tokenizer_class is not None:
+ tokenizer_class = fast_tokenizer_class
+ elif slow_tokenizer_class is not None:
+ tokenizer_class = slow_tokenizer_class
+ else:
+ raise ValueError("You need to pass a `tokenizer_class`")
+
+ for candidate in (slow_tokenizer_class, fast_tokenizer_class, tokenizer_class):
+ if candidate is not None:
+ REGISTERED_TOKENIZER_CLASSES[candidate.__name__] = candidate
+
+ if slow_tokenizer_class is not None and fast_tokenizer_class is not None:
+ REGISTERED_FAST_ALIASES[slow_tokenizer_class.__name__] = fast_tokenizer_class
+
+ TOKENIZER_MAPPING.register(config_class, tokenizer_class, exist_ok=exist_ok)
+
+
+__all__ = ["TOKENIZER_MAPPING", "AutoTokenizer"]
diff --git a/third_party/transformers/src/transformers/models/auto/video_processing_auto.py b/third_party/transformers/src/transformers/models/auto/video_processing_auto.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca6faa440ac8a3453c93ec33b7b4e561459f7ad9
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/auto/video_processing_auto.py
@@ -0,0 +1,424 @@
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# 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.
+"""AutoVideoProcessor class."""
+
+import importlib
+import os
+from collections import OrderedDict
+from typing import TYPE_CHECKING
+
+# Build the list of all video processors
+from ...configuration_utils import PreTrainedConfig
+from ...dynamic_module_utils import get_class_from_dynamic_module, resolve_trust_remote_code
+from ...utils import (
+ CONFIG_NAME,
+ IMAGE_PROCESSOR_NAME,
+ PROCESSOR_NAME,
+ VIDEO_PROCESSOR_NAME,
+ cached_file,
+ is_torchvision_available,
+ logging,
+ safe_load_json_file,
+)
+from ...utils.import_utils import requires
+from ...video_processing_utils import BaseVideoProcessor
+from .auto_factory import _LazyAutoMapping
+from .configuration_auto import (
+ CONFIG_MAPPING_NAMES,
+ AutoConfig,
+ model_type_to_module_name,
+ replace_list_option_in_docstrings,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+if TYPE_CHECKING:
+ # This significantly improves completion suggestion performance when
+ # the transformers package is used with Microsoft's Pylance language server.
+ VIDEO_PROCESSOR_MAPPING_NAMES: OrderedDict[str, tuple[str | None, str | None]] = OrderedDict()
+else:
+ VIDEO_PROCESSOR_MAPPING_NAMES = OrderedDict(
+ [
+ ("ernie4_5_vl_moe", "Ernie4_5_VLMoeVideoProcessor"),
+ ("gemma4", "Gemma4VideoProcessor"),
+ ("glm46v", "Glm46VVideoProcessor"),
+ ("glm4v", "Glm4vVideoProcessor"),
+ ("instructblip", "InstructBlipVideoVideoProcessor"),
+ ("instructblipvideo", "InstructBlipVideoVideoProcessor"),
+ ("internvl", "InternVLVideoProcessor"),
+ ("llava_next_video", "LlavaNextVideoVideoProcessor"),
+ ("llava_onevision", "LlavaOnevisionVideoProcessor"),
+ ("pe_audio_video", "PeVideoVideoProcessor"),
+ ("pe_video", "PeVideoVideoProcessor"),
+ ("perception_lm", "PerceptionLMVideoProcessor"),
+ ("qwen2_5_omni", "Qwen2VLVideoProcessor"),
+ ("qwen2_5_vl", "Qwen2VLVideoProcessor"),
+ ("qwen2_vl", "Qwen2VLVideoProcessor"),
+ ("qwen3_5", "Qwen3VLVideoProcessor"),
+ ("qwen3_5_moe", "Qwen3VLVideoProcessor"),
+ ("qwen3_omni_moe", "Qwen2VLVideoProcessor"),
+ ("qwen3_vl", "Qwen3VLVideoProcessor"),
+ ("qwen3_vl_moe", "Qwen3VLVideoProcessor"),
+ ("sam2_video", "Sam2VideoVideoProcessor"),
+ ("sam3_video", "Sam3VideoVideoProcessor"),
+ ("smolvlm", "SmolVLMVideoProcessor"),
+ ("video_llama_3", "VideoLlama3VideoProcessor"),
+ ("video_llava", "VideoLlavaVideoProcessor"),
+ ("videomae", "VideoMAEVideoProcessor"),
+ ("videomt", "VideomtVideoProcessor"),
+ ("vjepa2", "VJEPA2VideoProcessor"),
+ ]
+ )
+
+for model_type, video_processors in VIDEO_PROCESSOR_MAPPING_NAMES.items():
+ fast_video_processor_class = video_processors
+
+ # If the torchvision is not available, we set it to None
+ if not is_torchvision_available():
+ fast_video_processor_class = None
+
+ VIDEO_PROCESSOR_MAPPING_NAMES[model_type] = fast_video_processor_class
+
+VIDEO_PROCESSOR_MAPPING = _LazyAutoMapping(CONFIG_MAPPING_NAMES, VIDEO_PROCESSOR_MAPPING_NAMES)
+
+
+def video_processor_class_from_name(class_name: str):
+ for module_name, extractor in VIDEO_PROCESSOR_MAPPING_NAMES.items():
+ if class_name == extractor:
+ module_name = model_type_to_module_name(module_name)
+
+ module = importlib.import_module(f".{module_name}", "transformers.models")
+ try:
+ return getattr(module, class_name)
+ except AttributeError:
+ continue
+
+ for extractor in VIDEO_PROCESSOR_MAPPING._extra_content.values():
+ if getattr(extractor, "__name__", None) == class_name:
+ return extractor
+
+ # We did not find the class, but maybe it's because a dep is missing. In that case, the class will be in the main
+ # init and we return the proper dummy to get an appropriate error message.
+ main_module = importlib.import_module("transformers")
+ if hasattr(main_module, class_name):
+ return getattr(main_module, class_name)
+
+ return None
+
+
+def get_video_processor_config(
+ pretrained_model_name_or_path: str | os.PathLike,
+ cache_dir: str | os.PathLike | None = None,
+ force_download: bool = False,
+ proxies: dict[str, str] | None = None,
+ token: bool | str | None = None,
+ revision: str | None = None,
+ local_files_only: bool = False,
+ **kwargs,
+):
+ """
+ Loads the video processor configuration from a pretrained model video processor configuration.
+
+ Args:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained model configuration hosted inside a model repo on
+ huggingface.co.
+ - a path to a *directory* containing a configuration file saved using the
+ [`~BaseVideoProcessor.save_pretrained`] method, e.g., `./my_model_directory/`.
+
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model configuration should be cached if the standard
+ cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the configuration files and override the cached versions if they
+ exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `hf auth login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ local_files_only (`bool`, *optional*, defaults to `False`):
+ If `True`, will only try to load the video processor configuration from local files.
+
+
+
+ Passing `token=True` is required when you want to use a private model.
+
+
+
+ Returns:
+ `Dict`: The configuration of the video processor.
+
+ Examples:
+
+ ```python
+ # Download configuration from huggingface.co and cache.
+ video_processor_config = get_video_processor_config("llava-hf/llava-onevision-qwen2-0.5b-ov-hf")
+ # This model does not have a video processor config so the result will be an empty dict.
+ video_processor_config = get_video_processor_config("FacebookAI/xlm-roberta-base")
+
+ # Save a pretrained video processor locally and you can reload its config
+ from transformers import AutoVideoProcessor
+
+ video_processor = AutoVideoProcessor.from_pretrained("llava-hf/llava-onevision-qwen2-0.5b-ov-hf")
+ video_processor.save_pretrained("video-processor-test")
+ video_processor = get_video_processor_config("video-processor-test")
+ ```"""
+ # Load with a priority given to the nested processor config, if available in repo
+ resolved_processor_file = cached_file(
+ pretrained_model_name_or_path,
+ filename=PROCESSOR_NAME,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ token=token,
+ revision=revision,
+ local_files_only=local_files_only,
+ _raise_exceptions_for_gated_repo=False,
+ _raise_exceptions_for_missing_entries=False,
+ )
+ resolved_video_processor_files = [
+ resolved_file
+ for filename in [VIDEO_PROCESSOR_NAME, IMAGE_PROCESSOR_NAME]
+ if (
+ resolved_file := cached_file(
+ pretrained_model_name_or_path,
+ filename=filename,
+ cache_dir=cache_dir,
+ force_download=force_download,
+ proxies=proxies,
+ token=token,
+ revision=revision,
+ local_files_only=local_files_only,
+ _raise_exceptions_for_gated_repo=False,
+ _raise_exceptions_for_missing_entries=False,
+ _raise_exceptions_for_connection_errors=False,
+ )
+ )
+ is not None
+ ]
+ resolved_video_processor_file = resolved_video_processor_files[0] if resolved_video_processor_files else None
+
+ # An empty list if none of the possible files is found in the repo
+ if not resolved_video_processor_file and not resolved_processor_file:
+ logger.info("Could not locate the video processor configuration file.")
+ return {}
+
+ # Load video_processor dict. Priority goes as (nested config if found -> video processor config -> image processor config)
+ # We are downloading both configs because almost all models have a `processor_config.json` but
+ # not all of these are nested. We need to check if it was saved recebtly as nested or if it is legacy style
+ video_processor_dict = {}
+ if resolved_processor_file is not None:
+ processor_dict = safe_load_json_file(resolved_processor_file)
+ if "video_processor" in processor_dict:
+ video_processor_dict = processor_dict["video_processor"]
+
+ if resolved_video_processor_file is not None and video_processor_dict is None:
+ video_processor_dict = safe_load_json_file(resolved_video_processor_file)
+
+ return video_processor_dict
+
+
+@requires(backends=("vision", "torchvision"))
+class AutoVideoProcessor:
+ r"""
+ This is a generic video processor class that will be instantiated as one of the video processor classes of the
+ library when created with the [`AutoVideoProcessor.from_pretrained`] class method.
+
+ This class cannot be instantiated directly using `__init__()` (throws an error).
+ """
+
+ def __init__(self):
+ raise OSError(
+ "AutoVideoProcessor is designed to be instantiated "
+ "using the `AutoVideoProcessor.from_pretrained(pretrained_model_name_or_path)` method."
+ )
+
+ @classmethod
+ @replace_list_option_in_docstrings(VIDEO_PROCESSOR_MAPPING_NAMES)
+ def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
+ r"""
+ Instantiate one of the video processor classes of the library from a pretrained model vocabulary.
+
+ The video processor class to instantiate is selected based on the `model_type` property of the config object
+ (either passed as an argument or loaded from `pretrained_model_name_or_path` if possible), or when it's
+ missing, by falling back to using pattern matching on `pretrained_model_name_or_path`:
+
+ List options
+
+ Params:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained video_processor hosted inside a model repo on
+ huggingface.co.
+ - a path to a *directory* containing a video processor file saved using the
+ [`~video_processing_utils.BaseVideoProcessor.save_pretrained`] method, e.g.,
+ `./my_model_directory/`.
+ - a path to a saved video processor JSON *file*, e.g.,
+ `./my_model_directory/preprocessor_config.json`.
+ cache_dir (`str` or `os.PathLike`, *optional*):
+ Path to a directory in which a downloaded pretrained model video processor should be cached if the
+ standard cache should not be used.
+ force_download (`bool`, *optional*, defaults to `False`):
+ Whether or not to force to (re-)download the video processor files and override the cached versions if
+ they exist.
+ proxies (`dict[str, str]`, *optional*):
+ A dictionary of proxy servers to use by protocol or endpoint, e.g., `{'http': 'foo.bar:3128',
+ 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+ token (`str` or *bool*, *optional*):
+ The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
+ when running `hf auth login` (stored in `~/.huggingface`).
+ revision (`str`, *optional*, defaults to `"main"`):
+ The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+ git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
+ identifier allowed by git.
+ return_unused_kwargs (`bool`, *optional*, defaults to `False`):
+ If `False`, then this function returns just the final video processor object. If `True`, then this
+ functions returns a `Tuple(video_processor, unused_kwargs)` where *unused_kwargs* is a dictionary
+ consisting of the key/value pairs whose keys are not video processor attributes: i.e., the part of
+ `kwargs` which has not been used to update `video_processor` and is otherwise ignored.
+ trust_remote_code (`bool`, *optional*, defaults to `False`):
+ Whether or not to allow for custom models defined on the Hub in their own modeling files. This option
+ should only be set to `True` for repositories you trust and in which you have read the code, as it will
+ execute code present on the Hub on your local machine.
+ kwargs (`dict[str, Any]`, *optional*):
+ The values in kwargs of any keys which are video processor attributes will be used to override the
+ loaded values. Behavior concerning key/value pairs whose keys are *not* video processor attributes is
+ controlled by the `return_unused_kwargs` keyword parameter.
+
+
+
+ Passing `token=True` is required when you want to use a private model.
+
+
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoVideoProcessor
+
+ >>> # Download video processor from huggingface.co and cache.
+ >>> video_processor = AutoVideoProcessor.from_pretrained("llava-hf/llava-onevision-qwen2-0.5b-ov-hf")
+
+ >>> # If video processor files are in a directory (e.g. video processor was saved using *save_pretrained('./test/saved_model/')*)
+ >>> # video_processor = AutoVideoProcessor.from_pretrained("./test/saved_model/")
+ ```"""
+ config = kwargs.pop("config", None)
+ trust_remote_code = kwargs.pop("trust_remote_code", None)
+ kwargs["_from_auto"] = True
+
+ config_dict, _ = BaseVideoProcessor.get_video_processor_dict(pretrained_model_name_or_path, **kwargs)
+ video_processor_class = config_dict.get("video_processor_type", None)
+ video_processor_auto_map = None
+ if "AutoVideoProcessor" in config_dict.get("auto_map", {}):
+ video_processor_auto_map = config_dict["auto_map"]["AutoVideoProcessor"]
+
+ # If we still don't have the video processor class, check if we're loading from a previous image processor config
+ # and if so, infer the video processor class from there.
+ if video_processor_class is None and video_processor_auto_map is None:
+ image_processor_class = config_dict.pop("image_processor_type", None)
+ if image_processor_class is not None:
+ video_processor_class_inferred = image_processor_class.replace("ImageProcessor", "VideoProcessor")
+
+ # Some models have different image processors, e.g. InternVL uses GotOCRImageProcessor
+ # We cannot use GotOCRVideoProcessor when falling back for BC and should try to infer from config later on
+ if video_processor_class_from_name(video_processor_class_inferred) is not None:
+ video_processor_class = video_processor_class_inferred
+ if "AutoImageProcessor" in config_dict.get("auto_map", {}):
+ image_processor_auto_map = config_dict["auto_map"]["AutoImageProcessor"]
+ video_processor_auto_map = image_processor_auto_map.replace("ImageProcessor", "VideoProcessor")
+
+ # If we don't find the video processor class in the video processor config, let's try the model config.
+ if video_processor_class is None and video_processor_auto_map is None:
+ if not isinstance(config, PreTrainedConfig):
+ config = AutoConfig.from_pretrained(
+ pretrained_model_name_or_path, trust_remote_code=trust_remote_code, **kwargs
+ )
+ # It could be in `config.video_processor_type``
+ video_processor_class = getattr(config, "video_processor_type", None)
+ if hasattr(config, "auto_map") and "AutoVideoProcessor" in config.auto_map:
+ video_processor_auto_map = config.auto_map["AutoVideoProcessor"]
+
+ if video_processor_class is not None:
+ video_processor_class = video_processor_class_from_name(video_processor_class)
+
+ has_remote_code = video_processor_auto_map is not None
+ has_local_code = video_processor_class is not None or type(config) in VIDEO_PROCESSOR_MAPPING
+ explicit_local_code = has_local_code and not (
+ video_processor_class or VIDEO_PROCESSOR_MAPPING[type(config)]
+ ).__module__.startswith("transformers.")
+ if has_remote_code:
+ if "--" in video_processor_auto_map:
+ upstream_repo = video_processor_auto_map.split("--")[0]
+ else:
+ upstream_repo = None
+ trust_remote_code = resolve_trust_remote_code(
+ trust_remote_code, pretrained_model_name_or_path, has_local_code, has_remote_code, upstream_repo
+ )
+
+ if has_remote_code and trust_remote_code and not explicit_local_code:
+ class_ref = video_processor_auto_map
+ video_processor_class = get_class_from_dynamic_module(class_ref, pretrained_model_name_or_path, **kwargs)
+ _ = kwargs.pop("code_revision", None)
+ video_processor_class.register_for_auto_class()
+ return video_processor_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+ elif video_processor_class is not None:
+ return video_processor_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+ # Last try: we use the VIDEO_PROCESSOR_MAPPING.
+ elif type(config) in VIDEO_PROCESSOR_MAPPING:
+ video_processor_class = VIDEO_PROCESSOR_MAPPING[type(config)]
+ if video_processor_class is not None:
+ return video_processor_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)
+
+ # Raise a more informative error message if torchvision isn't found, otherwise just fallback to default
+ if not is_torchvision_available():
+ raise ValueError(
+ f"{pretrained_model_name_or_path} requires `torchvision` to be installed. Please install `torchvision` and try again."
+ )
+
+ raise ValueError(
+ f"Unrecognized video processor in {pretrained_model_name_or_path}. Should have a "
+ f"`video_processor_type` key in its {VIDEO_PROCESSOR_NAME} of {CONFIG_NAME}, or one of the following "
+ f"`model_type` keys in its {CONFIG_NAME}: {', '.join(c for c in VIDEO_PROCESSOR_MAPPING_NAMES)}"
+ )
+
+ @staticmethod
+ def register(
+ config_class,
+ video_processor_class,
+ exist_ok=False,
+ ):
+ """
+ Register a new video processor for this class.
+
+ Args:
+ config_class ([`PreTrainedConfig`]):
+ The configuration corresponding to the model to register.
+ video_processor_class ([`BaseVideoProcessor`]):
+ The video processor to register.
+ """
+ VIDEO_PROCESSOR_MAPPING.register(config_class, video_processor_class, exist_ok=exist_ok)
+
+
+__all__ = ["VIDEO_PROCESSOR_MAPPING", "AutoVideoProcessor"]
diff --git a/third_party/transformers/src/transformers/models/autoformer/__init__.py b/third_party/transformers/src/transformers/models/autoformer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..48a329608039b1a4a96ac1f99c20ee427f845cd9
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/autoformer/__init__.py
@@ -0,0 +1,27 @@
+# Copyright 2023 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_autoformer import *
+ from .modeling_autoformer import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/autoformer/configuration_autoformer.py b/third_party/transformers/src/transformers/models/autoformer/configuration_autoformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..5dc7d342ebcb8a25a0c27b85b62d43bcbd7cb058
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/autoformer/configuration_autoformer.py
@@ -0,0 +1,169 @@
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Autoformer model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="huggingface/autoformer-tourism-monthly")
+@strict
+class AutoformerConfig(PreTrainedConfig):
+ r"""
+ prediction_length (`int`):
+ The prediction length for the decoder. In other words, the prediction horizon of the model.
+ context_length (`int`, *optional*, defaults to `prediction_length`):
+ The context length for the encoder. If unset, the context length will be the same as the
+ `prediction_length`.
+ distribution_output (`string`, *optional*, defaults to `"student_t"`):
+ The distribution emission head for the model. Could be either "student_t", "normal" or "negative_binomial".
+ loss (`string`, *optional*, defaults to `"nll"`):
+ The loss function for the model corresponding to the `distribution_output` head. For parametric
+ distributions it is the negative log likelihood (nll) - which currently is the only supported one.
+ input_size (`int`, *optional*, defaults to 1):
+ The size of the target variable which by default is 1 for univariate targets. Would be > 1 in case of
+ multivariate targets.
+ lags_sequence (`list[int]`, *optional*, defaults to `[1, 2, 3, 4, 5, 6, 7]`):
+ The lags of the input time series as covariates often dictated by the frequency. Default is `[1, 2, 3, 4,
+ 5, 6, 7]`.
+ scaling (`bool`, *optional* defaults to `True`):
+ Whether to scale the input targets.
+ num_time_features (`int`, *optional*, defaults to 0):
+ The number of time features in the input time series.
+ num_dynamic_real_features (`int`, *optional*, defaults to 0):
+ The number of dynamic real valued features.
+ num_static_categorical_features (`int`, *optional*, defaults to 0):
+ The number of static categorical features.
+ num_static_real_features (`int`, *optional*, defaults to 0):
+ The number of static real valued features.
+ cardinality (`list[int]`, *optional*):
+ The cardinality (number of different values) for each of the static categorical features. Should be a list
+ of integers, having the same length as `num_static_categorical_features`. Cannot be `None` if
+ `num_static_categorical_features` is > 0.
+ num_parallel_samples (`int`, *optional*, defaults to 100):
+ The number of samples to generate in parallel for each time step of inference.
+ label_length (`int`, *optional*, defaults to 10):
+ Start token length of the Autoformer decoder, which is used for direct multi-step prediction (i.e.
+ non-autoregressive generation).
+ moving_average (`int`, *optional*, defaults to 25):
+ The window size of the moving average. In practice, it's the kernel size in AvgPool1d of the Decomposition
+ Layer.
+ autocorrelation_factor (`int`, *optional*, defaults to 3):
+ "Attention" (i.e. AutoCorrelation mechanism) factor which is used to find top k autocorrelations delays.
+ It's recommended in the paper to set it to a number between 1 and 5.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoformerConfig, AutoformerModel
+
+ >>> # Initializing a default Autoformer configuration
+ >>> configuration = AutoformerConfig()
+
+ >>> # Randomly initializing a model (with random weights) from the configuration
+ >>> model = AutoformerModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "autoformer"
+ attribute_map = {
+ "hidden_size": "d_model",
+ "num_attention_heads": "encoder_attention_heads",
+ "num_hidden_layers": "encoder_layers",
+ }
+
+ prediction_length: int | None = None
+ context_length: int | None = None
+ distribution_output: str = "student_t"
+ loss: str = "nll"
+ input_size: int = 1
+ lags_sequence: list[int] | tuple[int, ...] = (1, 2, 3, 4, 5, 6, 7)
+ scaling: bool | str = True
+ num_time_features: int = 0
+ num_dynamic_real_features: int = 0
+ num_static_categorical_features: int = 0
+ num_static_real_features: int = 0
+ cardinality: list[int] | None = None
+ embedding_dimension: list[int] | None = None
+ d_model: int = 64
+ encoder_attention_heads: int = 2
+ decoder_attention_heads: int = 2
+ encoder_layers: int = 2
+ decoder_layers: int = 2
+ encoder_ffn_dim: int = 32
+ decoder_ffn_dim: int = 32
+ activation_function: str = "gelu"
+ dropout: float | int = 0.1
+ encoder_layerdrop: float | int = 0.1
+ decoder_layerdrop: float | int = 0.1
+ attention_dropout: float | int = 0.1
+ activation_dropout: float | int = 0.1
+ num_parallel_samples: int = 100
+ init_std: float = 0.02
+ use_cache: bool = True
+ is_encoder_decoder: bool = True
+ label_length: int = 10
+ moving_average: int = 25
+ autocorrelation_factor: int = 3
+
+ def __post_init__(self, **kwargs):
+ self.context_length = self.context_length if self.context_length is not None else self.prediction_length
+ self.lags_sequence = list(self.lags_sequence)
+
+ if not (self.cardinality is not None and self.num_static_categorical_features > 0):
+ self.cardinality = [0]
+
+ if not (self.embedding_dimension is not None and self.num_static_categorical_features > 0):
+ self.embedding_dimension = [min(50, (cat + 1) // 2) for cat in self.cardinality]
+
+ # Transformer architecture configuration
+ self.feature_size = self.input_size * len(self.lags_sequence) + self._number_of_features
+ super().__post_init__(**kwargs)
+
+ @property
+ def _number_of_features(self) -> int:
+ return (
+ sum(self.embedding_dimension)
+ + self.num_dynamic_real_features
+ + self.num_time_features
+ + self.num_static_real_features
+ + self.input_size * 2 # the log1p(abs(loc)) and log(scale) features
+ )
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ if (
+ self.cardinality is not None
+ and self.num_static_categorical_features > 0
+ and len(self.cardinality) != self.num_static_categorical_features
+ ):
+ raise ValueError(
+ "The cardinality should be a list of the same length as `num_static_categorical_features`"
+ )
+
+ if (
+ self.embedding_dimension is not None
+ and self.num_static_categorical_features > 0
+ and len(self.embedding_dimension) != self.num_static_categorical_features
+ ):
+ raise ValueError(
+ "The embedding dimension should be a list of the same length as `num_static_categorical_features`"
+ )
+
+
+__all__ = ["AutoformerConfig"]
diff --git a/third_party/transformers/src/transformers/models/autoformer/modeling_autoformer.py b/third_party/transformers/src/transformers/models/autoformer/modeling_autoformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..9ff0403b52f605de17935e441b061518f5efaf77
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/autoformer/modeling_autoformer.py
@@ -0,0 +1,1832 @@
+# Copyright (c) 2021 THUML @ Tsinghua University
+# Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# 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 Autoformer model."""
+
+import math
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...masking_utils import create_bidirectional_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutput, ModelOutput, SampleTSPredictionOutput, Seq2SeqTSPredictionOutput
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...time_series_utils import NegativeBinomialOutput, NormalOutput, StudentTOutput
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from .configuration_autoformer import AutoformerConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding).
+ """
+)
+class AutoFormerDecoderOutput(ModelOutput):
+ r"""
+ 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.
+ trend (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ Trend tensor for each time series.
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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.
+ cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=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 of the decoder's cross-attention layer, after the attention softmax, used to compute the
+ weighted average in the cross-attention heads.
+ """
+
+ last_hidden_state: torch.FloatTensor | None = None
+ trend: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ cross_attentions: tuple[torch.FloatTensor] | None = None
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Autoformer model output that contains the additional trend output.
+ """
+)
+class AutoformerModelOutput(ModelOutput):
+ r"""
+ 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 decoder 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.
+ trend (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ Trend tensor for each time series.
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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.
+ loc (`torch.FloatTensor` of shape `(batch_size,)` or `(batch_size, input_size)`, *optional*):
+ Shift values of each time series' context window which is used to give the model inputs of the same
+ magnitude and then used to shift back to the original magnitude.
+ scale (`torch.FloatTensor` of shape `(batch_size,)` or `(batch_size, input_size)`, *optional*):
+ Scaling values of each time series' context window which is used to give the model inputs of the same
+ magnitude and then used to rescale back to the original magnitude.
+ static_features: (`torch.FloatTensor` of shape `(batch_size, feature size)`, *optional*):
+ Static features of each time series' in a batch which are copied to the covariates at inference time.
+ """
+
+ last_hidden_state: torch.FloatTensor | None = None
+ trend: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ decoder_hidden_states: tuple[torch.FloatTensor] | None = None
+ decoder_attentions: tuple[torch.FloatTensor] | None = None
+ cross_attentions: tuple[torch.FloatTensor] | None = None
+ encoder_last_hidden_state: torch.FloatTensor | None = None
+ encoder_hidden_states: tuple[torch.FloatTensor] | None = None
+ encoder_attentions: tuple[torch.FloatTensor] | None = None
+ loc: torch.FloatTensor | None = None
+ scale: torch.FloatTensor | None = None
+ static_features: torch.FloatTensor | None = None
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesFeatureEmbedder with TimeSeries->Autoformer
+class AutoformerFeatureEmbedder(nn.Module):
+ """
+ Embed a sequence of categorical features.
+
+ Args:
+ cardinalities (`list[int]`):
+ List of cardinalities of the categorical features.
+ embedding_dims (`list[int]`):
+ List of embedding dimensions of the categorical features.
+ """
+
+ def __init__(self, cardinalities: list[int], embedding_dims: list[int]) -> None:
+ super().__init__()
+
+ self.num_features = len(cardinalities)
+ self.embedders = nn.ModuleList([nn.Embedding(c, d) for c, d in zip(cardinalities, embedding_dims)])
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ if self.num_features > 1:
+ # we slice the last dimension, giving an array of length
+ # self.num_features with shape (N,T) or (N)
+ cat_feature_slices = torch.chunk(features, self.num_features, dim=-1)
+ else:
+ cat_feature_slices = [features]
+
+ return torch.cat(
+ [
+ embed(cat_feature_slice.squeeze(-1))
+ for embed, cat_feature_slice in zip(self.embedders, cat_feature_slices)
+ ],
+ dim=-1,
+ )
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesStdScaler with TimeSeriesTransformer->Autoformer,TimeSeries->Autoformer
+class AutoformerStdScaler(nn.Module):
+ """
+ Standardize features by calculating the mean and scaling along the first dimension, and then normalizes it by
+ subtracting from the mean and dividing by the standard deviation.
+ """
+
+ def __init__(self, config: AutoformerConfig):
+ super().__init__()
+ self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1
+ self.keepdim = config.keepdim if hasattr(config, "keepdim") else True
+ self.minimum_scale = config.minimum_scale if hasattr(config, "minimum_scale") else 1e-5
+
+ def forward(
+ self, data: torch.Tensor, observed_indicator: torch.Tensor
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Parameters:
+ data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+ input for Batch norm calculation
+ observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+ Calculating the scale on the observed indicator.
+ Returns:
+ tuple of `torch.Tensor` of shapes
+ (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,
+ `(batch_size, 1, num_input_channels)`)
+ """
+ denominator = observed_indicator.sum(self.dim, keepdim=self.keepdim)
+ denominator = denominator.clamp_min(1.0)
+ loc = (data * observed_indicator).sum(self.dim, keepdim=self.keepdim) / denominator
+
+ variance = (((data - loc) * observed_indicator) ** 2).sum(self.dim, keepdim=self.keepdim) / denominator
+ scale = torch.sqrt(variance + self.minimum_scale)
+ return (data - loc) / scale, loc, scale
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesMeanScaler with TimeSeriesTransformer->Autoformer,TimeSeries->Autoformer
+class AutoformerMeanScaler(nn.Module):
+ """
+ Computes a scaling factor as the weighted average absolute value along the first dimension, and scales the data
+ accordingly.
+ """
+
+ def __init__(self, config: AutoformerConfig):
+ super().__init__()
+ self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1
+ self.keepdim = config.keepdim if hasattr(config, "keepdim") else True
+ self.minimum_scale = config.minimum_scale if hasattr(config, "minimum_scale") else 1e-10
+ self.default_scale = config.default_scale if hasattr(config, "default_scale") else None
+
+ def forward(
+ self, data: torch.Tensor, observed_indicator: torch.Tensor
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Parameters:
+ data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+ input for Batch norm calculation
+ observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+ Calculating the scale on the observed indicator.
+ Returns:
+ tuple of `torch.Tensor` of shapes
+ (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,
+ `(batch_size, 1, num_input_channels)`)
+ """
+ ts_sum = (data * observed_indicator).abs().sum(self.dim, keepdim=True)
+ num_observed = observed_indicator.sum(self.dim, keepdim=True)
+
+ scale = ts_sum / torch.clamp(num_observed, min=1)
+
+ # If `default_scale` is provided, we use it, otherwise we use the scale
+ # of the batch.
+ if self.default_scale is None:
+ batch_sum = ts_sum.sum(dim=0)
+ batch_observations = torch.clamp(num_observed.sum(0), min=1)
+ default_scale = torch.squeeze(batch_sum / batch_observations)
+ else:
+ default_scale = self.default_scale * torch.ones_like(scale)
+
+ # apply default scale where there are no observations
+ scale = torch.where(num_observed > 0, scale, default_scale)
+
+ # ensure the scale is at least `self.minimum_scale`
+ scale = torch.clamp(scale, min=self.minimum_scale)
+ scaled_data = data / scale
+
+ if not self.keepdim:
+ scale = scale.squeeze(dim=self.dim)
+
+ return scaled_data, torch.zeros_like(scale), scale
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesNOPScaler with TimeSeriesTransformer->Autoformer,TimeSeries->Autoformer
+class AutoformerNOPScaler(nn.Module):
+ """
+ Assigns a scaling factor equal to 1 along the first dimension, and therefore applies no scaling to the input data.
+ """
+
+ def __init__(self, config: AutoformerConfig):
+ super().__init__()
+ self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1
+ self.keepdim = config.keepdim if hasattr(config, "keepdim") else True
+
+ def forward(
+ self, data: torch.Tensor, observed_indicator: torch.Tensor | None = None
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Parameters:
+ data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+ input for Batch norm calculation
+ Returns:
+ tuple of `torch.Tensor` of shapes
+ (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,
+ `(batch_size, 1, num_input_channels)`)
+ """
+ scale = torch.ones_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim)
+ loc = torch.zeros_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim)
+ return data, loc, scale
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.weighted_average
+def weighted_average(input_tensor: torch.Tensor, weights: torch.Tensor | None = None, dim=None) -> torch.Tensor:
+ """
+ Computes the weighted average of a given tensor across a given `dim`, masking values associated with weight zero,
+ meaning instead of `nan * 0 = nan` you will get `0 * 0 = 0`.
+
+ Args:
+ input_tensor (`torch.FloatTensor`):
+ Input tensor, of which the average must be computed.
+ weights (`torch.FloatTensor`, *optional*):
+ Weights tensor, of the same shape as `input_tensor`.
+ dim (`int`, *optional*):
+ The dim along which to average `input_tensor`.
+
+ Returns:
+ `torch.FloatTensor`: The tensor with values averaged along the specified `dim`.
+ """
+ if weights is not None:
+ weighted_tensor = torch.where(weights != 0, input_tensor * weights, torch.zeros_like(input_tensor))
+ sum_weights = torch.clamp(weights.sum(dim=dim) if dim else weights.sum(), min=1.0)
+ return (weighted_tensor.sum(dim=dim) if dim else weighted_tensor.sum()) / sum_weights
+ else:
+ return input_tensor.mean(dim=dim)
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.nll
+def nll(input: torch.distributions.Distribution, target: torch.Tensor) -> torch.Tensor:
+ """
+ Computes the negative log likelihood loss from input distribution with respect to target.
+ """
+ return -input.log_prob(target)
+
+
+# Copied from transformers.models.marian.modeling_marian.MarianSinusoidalPositionalEmbedding with Marian->Autoformer
+class AutoformerSinusoidalPositionalEmbedding(nn.Embedding):
+ """This module produces sinusoidal positional embeddings of any length."""
+
+ def __init__(self, num_positions: int, embedding_dim: int, padding_idx: int | None = None) -> None:
+ super().__init__(num_positions, embedding_dim, _freeze=True)
+
+ def create_weight(self):
+ """
+ Identical to the XLM create_sinusoidal_embeddings except features are not interleaved. The cos features are in
+ the 2nd half of the vector. [dim // 2:]
+ """
+ n_pos, dim = self.weight.shape
+ position_enc = np.array(
+ [[pos / np.power(10000, 2 * (j // 2) / dim) for j in range(dim)] for pos in range(n_pos)]
+ )
+ out = torch.empty(n_pos, dim, dtype=self.weight.dtype, requires_grad=False)
+ sentinel = dim // 2 if dim % 2 == 0 else (dim // 2) + 1
+ out[:, 0:sentinel] = torch.FloatTensor(np.sin(position_enc[:, 0::2]))
+ out[:, sentinel:] = torch.FloatTensor(np.cos(position_enc[:, 1::2]))
+ return out
+
+ @torch.no_grad()
+ def forward(
+ self, input_ids_shape: torch.Size, past_key_values_length: int = 0, position_ids: torch.Tensor | None = None
+ ) -> torch.Tensor:
+ """`input_ids_shape` is expected to be [bsz x seqlen]."""
+ if position_ids is None:
+ bsz, seq_len = input_ids_shape[:2]
+ position_ids = torch.arange(
+ past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device
+ )
+ return super().forward(position_ids)
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesValueEmbedding with TimeSeries->Autoformer
+class AutoformerValueEmbedding(nn.Module):
+ def __init__(self, feature_size, d_model):
+ super().__init__()
+ self.value_projection = nn.Linear(in_features=feature_size, out_features=d_model, bias=False)
+
+ def forward(self, x):
+ return self.value_projection(x)
+
+
+# Class based on
+# https://github.com/thuml/Autoformer/blob/c6a0694ff484753f2d986cc0bb1f99ee850fc1a8/layers/Autoformer_EncDec.py#L39
+# where AutoformerSeriesDecompositionLayer is series_decomp + moving_average
+class AutoformerSeriesDecompositionLayer(nn.Module):
+ """
+ Returns the trend and the seasonal parts of the time series. Calculated as:
+
+ x_trend = AvgPool(Padding(X)) and x_seasonal = X - x_trend
+ """
+
+ def __init__(self, config: AutoformerConfig):
+ super().__init__()
+ self.kernel_size = config.moving_average
+ self.avg = nn.AvgPool1d(kernel_size=self.kernel_size, stride=1, padding=0)
+
+ def forward(self, x):
+ """Input shape: Batch x Time x EMBED_DIM"""
+ # padding on the both ends of time series
+ num_of_pads = (self.kernel_size - 1) // 2
+ front = x[:, 0:1, :].repeat(1, num_of_pads, 1)
+ end = x[:, -1:, :].repeat(1, num_of_pads, 1)
+ x_padded = torch.cat([front, x, end], dim=1)
+
+ # calculate the trend and seasonal part of the series
+ x_trend = self.avg(x_padded.permute(0, 2, 1)).permute(0, 2, 1)
+ x_seasonal = x - x_trend
+ return x_seasonal, x_trend
+
+
+# Class based on
+# https://github.com/thuml/Autoformer/blob/c6a0694ff484753f2d986cc0bb1f99ee850fc1a8/layers/Autoformer_EncDec.py#L6
+# where AutoformerLayernorm is my_Layernorm
+class AutoformerLayernorm(nn.Module):
+ """
+ Special designed layer normalization for the seasonal part, calculated as: AutoformerLayernorm(x) = nn.LayerNorm(x)
+ - torch.mean(nn.LayerNorm(x))
+ """
+
+ def __init__(self, config: AutoformerConfig):
+ super().__init__()
+ self.layernorm = nn.LayerNorm(config.d_model)
+
+ def forward(self, x):
+ x_hat = self.layernorm(x)
+ bias = torch.mean(x_hat, dim=1).unsqueeze(1).repeat(1, x.shape[1], 1)
+ return x_hat - bias
+
+
+class AutoformerAttention(nn.Module):
+ """
+ AutoCorrelation Mechanism with the following two phases:
+ (1) period-based dependencies discovery (2) time delay aggregation
+ This block replace the canonical self-attention mechanism.
+ """
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ dropout: float | None = 0.0,
+ is_decoder: bool | None = False,
+ bias: bool | None = True,
+ autocorrelation_factor: int | None = 3,
+ layer_idx: int | None = None,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.dropout = dropout
+ self.head_dim = embed_dim // num_heads
+
+ if (self.head_dim * num_heads) != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+ f" and `num_heads`: {num_heads})."
+ )
+ self.scaling = self.head_dim**-0.5
+ self.is_decoder = is_decoder
+ self.layer_idx = layer_idx
+
+ self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+ self.autocorrelation_factor = autocorrelation_factor
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ key_value_states: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Input shape: Batch x Time x Channel"""
+
+ # if key_value_states are provided this layer is used as a cross-attention layer
+ # for the decoder
+ is_cross_attention = key_value_states is not None
+ bsz, tgt_len, _ = hidden_states.size()
+
+ # get query proj
+ query_states = self.q_proj(hidden_states)
+
+ is_updated = False
+ if past_key_values is not None:
+ if isinstance(past_key_values, EncoderDecoderCache):
+ is_updated = past_key_values.is_updated.get(self.layer_idx)
+ if is_cross_attention:
+ # after the first generated id, we can subsequently re-use all key/value_layer from cache
+ curr_past_key_values = past_key_values.cross_attention_cache
+ else:
+ curr_past_key_values = past_key_values.self_attention_cache
+ else:
+ curr_past_key_values = past_key_values
+
+ current_states = key_value_states if is_cross_attention else hidden_states
+ if is_cross_attention and past_key_values is not None and is_updated:
+ # reuse k,v, cross_attentions
+ key_states = curr_past_key_values.layers[self.layer_idx].keys
+ value_states = curr_past_key_values.layers[self.layer_idx].values
+ else:
+ key_states = self.k_proj(current_states)
+ value_states = self.v_proj(current_states)
+ key_states = key_states.view(bsz, -1, self.num_heads, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, -1, self.num_heads, self.head_dim).transpose(1, 2)
+
+ if past_key_values is not None:
+ key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx)
+ # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+ if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):
+ past_key_values.is_updated[self.layer_idx] = True
+
+ proj_shape = (bsz * self.num_heads, -1, self.head_dim)
+ query_states = query_states.view(bsz, tgt_len, self.num_heads, self.head_dim).transpose(1, 2)
+ query_states = query_states.reshape(*proj_shape)
+ key_states = key_states.reshape(*proj_shape)
+ value_states = value_states.reshape(*proj_shape)
+
+ # (1) period-based dependencies discovery
+ # Resize (truncation or zero filling)
+ queries_time_length = query_states.size(1)
+ values_time_length = value_states.size(1)
+ if queries_time_length > values_time_length:
+ query_states = query_states[:, : (queries_time_length - values_time_length), :]
+ zeros = torch.zeros_like(query_states).float()
+ value_states = torch.cat([value_states, zeros], dim=1)
+ key_states = torch.cat([key_states, zeros], dim=1)
+ else:
+ value_states = value_states[:, :queries_time_length, :]
+ key_states = key_states[:, :queries_time_length, :]
+
+ query_states_fft = torch.fft.rfft(query_states, n=tgt_len, dim=1)
+ key_states_fft = torch.fft.rfft(key_states, n=tgt_len, dim=1)
+ attn_weights = query_states_fft * torch.conj(key_states_fft)
+ attn_weights = torch.fft.irfft(attn_weights, n=tgt_len, dim=1) # Autocorrelation(Q,K)
+
+ src_len = key_states.size(1)
+ channel = key_states.size(2)
+
+ if attn_weights.size() != (bsz * self.num_heads, tgt_len, channel):
+ raise ValueError(
+ f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, channel)}, but is"
+ f" {attn_weights.size()}"
+ )
+
+ if attention_mask is not None:
+ if attention_mask.size() != (bsz, 1, tgt_len, src_len):
+ raise ValueError(
+ f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
+ )
+ attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
+ attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+ # this operation is a bit awkward, but it's required to
+ # make sure that attn_weights keeps its gradient.
+ # In order to do so, attn_weights have to be reshaped
+ # twice and have to be reused in the following
+ attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, channel)
+ attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, channel)
+
+ # time delay aggregation
+ time_length = value_states.size(1)
+ autocorrelations = attn_weights.view(bsz, self.num_heads, tgt_len, channel)
+
+ # find top k autocorrelations delays
+ top_k = int(self.autocorrelation_factor * math.log(time_length))
+ autocorrelations_mean_on_head_channel = torch.mean(autocorrelations, dim=(1, -1)) # bsz x tgt_len
+ if self.training:
+ autocorrelations_mean_on_bsz = torch.mean(autocorrelations_mean_on_head_channel, dim=0)
+ _, top_k_delays_index = torch.topk(autocorrelations_mean_on_bsz, top_k)
+ top_k_autocorrelations = torch.stack(
+ [autocorrelations_mean_on_head_channel[:, top_k_delays_index[i]] for i in range(top_k)], dim=-1
+ )
+ else:
+ top_k_autocorrelations, top_k_delays_index = torch.topk(
+ autocorrelations_mean_on_head_channel, top_k, dim=1
+ )
+
+ top_k_autocorrelations = torch.softmax(top_k_autocorrelations, dim=-1) # bsz x top_k
+
+ # compute aggregation: value_states.roll(delay) * top_k_autocorrelations(delay)
+ if not self.training:
+ # used for compute values_states.roll(delay) in inference
+ tmp_values = value_states.repeat(1, 2, 1)
+ init_index = (
+ torch.arange(time_length)
+ .view(1, -1, 1)
+ .repeat(bsz * self.num_heads, 1, channel)
+ .to(value_states.device)
+ )
+
+ delays_agg = torch.zeros_like(value_states).float() # bsz x time_length x channel
+ for i in range(top_k):
+ # compute value_states roll delay
+ if not self.training:
+ tmp_delay = init_index + top_k_delays_index[:, i].view(-1, 1, 1).repeat(
+ self.num_heads, tgt_len, channel
+ )
+ value_states_roll_delay = torch.gather(tmp_values, dim=1, index=tmp_delay)
+ else:
+ value_states_roll_delay = value_states.roll(shifts=-int(top_k_delays_index[i]), dims=1)
+
+ # aggregation
+ top_k_autocorrelations_at_delay = (
+ top_k_autocorrelations[:, i].view(-1, 1, 1).repeat(self.num_heads, tgt_len, channel)
+ )
+ delays_agg += value_states_roll_delay * top_k_autocorrelations_at_delay
+
+ attn_output = delays_agg.contiguous()
+
+ if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
+ raise ValueError(
+ f"`attn_output` should be of size {(bsz * self.num_heads, tgt_len, self.head_dim)}, but is"
+ f" {attn_output.size()}"
+ )
+
+ attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
+ attn_output = attn_output.transpose(1, 2)
+
+ # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
+ # partitioned across GPUs when using tensor-parallelism.
+ attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)
+
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output, attn_weights_reshaped
+
+
+class AutoformerEncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: AutoformerConfig):
+ super().__init__()
+ self.embed_dim = config.d_model
+ self.self_attn = AutoformerAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.encoder_attention_heads,
+ dropout=config.attention_dropout,
+ autocorrelation_factor=config.autocorrelation_factor,
+ )
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.dropout = config.dropout
+ self.activation_fn = ACT2FN[config.activation_function]
+ self.activation_dropout = config.activation_dropout
+ self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
+ self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
+ self.final_layer_norm = AutoformerLayernorm(config)
+ self.decomp1 = AutoformerSeriesDecompositionLayer(config)
+ self.decomp2 = AutoformerSeriesDecompositionLayer(config)
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor,
+ attention_mask: torch.FloatTensor = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.FloatTensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ """
+ residual = hidden_states
+ hidden_states, _ = self.self_attn(
+ hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ # added layer norm here as an improvement
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+ hidden_states, _ = self.decomp1(hidden_states)
+
+ residual = hidden_states
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states, _ = self.decomp2(hidden_states)
+ hidden_states = self.final_layer_norm(hidden_states)
+
+ if hidden_states.dtype == torch.float16 and not torch.isfinite(hidden_states).all():
+ clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+ hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+ return hidden_states
+
+
+class AutoformerDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: AutoformerConfig, layer_idx=None):
+ super().__init__()
+ self.embed_dim = config.d_model
+
+ self.self_attn = AutoformerAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.decoder_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ autocorrelation_factor=config.autocorrelation_factor,
+ layer_idx=layer_idx,
+ )
+ self.dropout = config.dropout
+ self.activation_fn = ACT2FN[config.activation_function]
+ self.activation_dropout = config.activation_dropout
+
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.encoder_attn = AutoformerAttention(
+ self.embed_dim,
+ config.decoder_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ autocorrelation_factor=config.autocorrelation_factor,
+ layer_idx=layer_idx,
+ )
+ self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
+ self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
+ self.final_layer_norm = AutoformerLayernorm(config)
+
+ self.decomp1 = AutoformerSeriesDecompositionLayer(config)
+ self.decomp2 = AutoformerSeriesDecompositionLayer(config)
+ self.decomp3 = AutoformerSeriesDecompositionLayer(config)
+
+ # source: https://github.com/thuml/Autoformer/blob/e6371e24f2ae2dd53e472edefdd5814c5176f864/layers/Autoformer_EncDec.py#L128
+ self.trend_projection = nn.Conv1d(
+ in_channels=self.embed_dim,
+ out_channels=config.feature_size,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ padding_mode="circular",
+ bias=False,
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = True,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.FloatTensor, torch.FloatTensor]:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ encoder_hidden_states (`torch.FloatTensor`):
+ cross attention input to the layer of shape `(batch, seq_len, embed_dim)`
+ encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ past_key_values (`Cache`): cached past key and value projection states
+ use_cache: (`bool`, *optional*, defaults to `True`):
+ Whether or not the model should return the `present_key_value` state to be used for subsequent
+ decoding.
+ """
+ residual = hidden_states
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states, trend1 = self.decomp1(hidden_states)
+ # added layer norm here as an improvement
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ # Cross-Attention Block
+ if encoder_hidden_states is not None:
+ residual = hidden_states
+
+ hidden_states, _ = self.encoder_attn(
+ hidden_states,
+ key_value_states=encoder_hidden_states,
+ attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states, trend2 = self.decomp2(hidden_states)
+ # added layer norm here as an improvement
+ hidden_states = self.encoder_attn_layer_norm(hidden_states)
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states, trend3 = self.decomp3(hidden_states)
+ hidden_states = self.final_layer_norm(hidden_states)
+
+ if encoder_hidden_states is not None:
+ residual_trend = trend1 + trend2 + trend3
+ else:
+ residual_trend = trend1 + trend3
+ residual_trend = self.trend_projection(residual_trend.permute(0, 2, 1)).transpose(1, 2)
+
+ return hidden_states, residual_trend
+
+
+@auto_docstring
+class AutoformerPreTrainedModel(PreTrainedModel):
+ config: AutoformerConfig
+ base_model_prefix = "model"
+ input_modalities = ("time",)
+ main_input_name = "past_values"
+ supports_gradient_checkpointing = True
+ _can_record_outputs = {
+ "hidden_states": (AutoformerEncoderLayer, AutoformerDecoderLayer),
+ "attentions": AutoformerAttention,
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module: nn.Module):
+ std = self.config.init_std
+ if isinstance(module, (nn.Linear, nn.Conv1d)):
+ init.normal_(module.weight, mean=0.0, std=std)
+ if module.bias is not None:
+ init.zeros_(module.bias)
+ elif isinstance(module, AutoformerSinusoidalPositionalEmbedding):
+ init.copy_(module.weight, module.create_weight())
+ elif isinstance(module, nn.Embedding):
+ init.normal_(module.weight, mean=0.0, std=std)
+ # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag
+ if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False):
+ init.zeros_(module.weight[module.padding_idx])
+ elif isinstance(module, nn.LayerNorm):
+ init.ones_(module.weight)
+ init.zeros_(module.bias)
+
+
+# copied from transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesTransformerEncoder with TimeSeriesTransformer->Autoformer,TimeSeries->Autoformer
+class AutoformerEncoder(AutoformerPreTrainedModel):
+ """
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
+ [`AutoformerEncoderLayer`].
+
+ Args:
+ config: AutoformerConfig
+ """
+
+ def __init__(self, config: AutoformerConfig):
+ super().__init__(config)
+
+ self.dropout = config.dropout
+ self.layerdrop = config.encoder_layerdrop
+ if config.prediction_length is None:
+ raise ValueError("The `prediction_length` config needs to be specified.")
+
+ self.value_embedding = AutoformerValueEmbedding(feature_size=config.feature_size, d_model=config.d_model)
+ self.embed_positions = AutoformerSinusoidalPositionalEmbedding(
+ config.context_length + config.prediction_length, config.d_model
+ )
+ self.layers = nn.ModuleList([AutoformerEncoderLayer(config) for _ in range(config.encoder_layers)])
+ self.layernorm_embedding = nn.LayerNorm(config.d_model)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(
+ self,
+ attention_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutput:
+ r"""
+ Args:
+ 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)
+ 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.
+ """
+ hidden_states = self.value_embedding(inputs_embeds)
+ embed_pos = self.embed_positions(inputs_embeds.size())
+
+ hidden_states = self.layernorm_embedding(hidden_states + embed_pos)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ )
+
+ for idx, encoder_layer in enumerate(self.layers):
+ # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+ to_drop = False
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop: # skip the layer
+ to_drop = True
+
+ if not to_drop:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+class AutoformerDecoder(AutoformerPreTrainedModel):
+ """
+ Transformer decoder consisting of `config.decoder_layers` layers. Each layer is a [`AutoformerDecoderLayer`]
+
+ Args:
+ config: AutoformerConfig
+ """
+
+ _can_record_outputs = {
+ "hidden_states": AutoformerDecoderLayer,
+ "attentions": OutputRecorder(AutoformerAttention, index=1, layer_name="self_attn"),
+ "cross_attentions": OutputRecorder(AutoformerAttention, index=1, layer_name="encoder_attn"),
+ }
+
+ def __init__(self, config: AutoformerConfig):
+ super().__init__(config)
+ self.dropout = config.dropout
+ self.layerdrop = config.decoder_layerdrop
+ if config.prediction_length is None:
+ raise ValueError("The `prediction_length` config needs to be specified.")
+
+ self.value_embedding = AutoformerValueEmbedding(feature_size=config.feature_size, d_model=config.d_model)
+ self.embed_positions = AutoformerSinusoidalPositionalEmbedding(
+ config.context_length + config.prediction_length, config.d_model
+ )
+ self.layers = nn.ModuleList(
+ [AutoformerDecoderLayer(config, layer_idx=i) for i in range(config.decoder_layers)]
+ )
+ self.layernorm_embedding = nn.LayerNorm(config.d_model)
+
+ # https://github.com/thuml/Autoformer/blob/e6371e24f2ae2dd53e472edefdd5814c5176f864/models/Autoformer.py#L74
+ self.seasonality_projection = nn.Linear(config.d_model, config.feature_size)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(
+ self,
+ trend: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | AutoFormerDecoderOutput:
+ r"""
+ Args:
+ trend (`torch.FloatTensor` of shape `(batch_size, prediction_length, feature_size)`, *optional*):
+ The trend sequence to be fed to the decoder.
+ 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)
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
+ of the decoder.
+ encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):
+ Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. 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)
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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 `use_cache` is True, `past_key_values` key value states are returned and can be used to speed up
+ decoding (see `past_key_values`).
+ """
+ if use_cache and past_key_values is None:
+ past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+
+ encoder_attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=encoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ )
+
+ hidden_states = self.value_embedding(inputs_embeds)
+ embed_pos = self.embed_positions(
+ inputs_embeds.size(), past_key_values_length=self.config.context_length - self.config.label_length
+ )
+ hidden_states = self.layernorm_embedding(hidden_states + embed_pos)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ # decoder layers
+ for idx, decoder_layer in enumerate(self.layers):
+ # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop:
+ continue
+
+ hidden_states, residual_trend = decoder_layer(
+ hidden_states,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+ trend = trend + residual_trend
+
+ # project seasonality representation
+ hidden_states = self.seasonality_projection(hidden_states)
+
+ return AutoFormerDecoderOutput(
+ last_hidden_state=hidden_states,
+ trend=trend,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class AutoformerModel(AutoformerPreTrainedModel):
+ def __init__(self, config: AutoformerConfig):
+ super().__init__(config)
+
+ if config.scaling == "mean" or config.scaling is True:
+ self.scaler = AutoformerMeanScaler(config)
+ elif config.scaling == "std":
+ self.scaler = AutoformerStdScaler(config)
+ else:
+ self.scaler = AutoformerNOPScaler(config)
+
+ if config.num_static_categorical_features > 0:
+ self.embedder = AutoformerFeatureEmbedder(
+ cardinalities=config.cardinality, embedding_dims=config.embedding_dimension
+ )
+
+ # transformer encoder-decoder and mask initializer
+ self.encoder = AutoformerEncoder(config)
+ self.decoder = AutoformerDecoder(config)
+
+ # used for decoder seasonal and trend initialization
+ self.decomposition_layer = AutoformerSeriesDecompositionLayer(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @property
+ def _past_length(self) -> int:
+ return self.config.context_length + max(self.config.lags_sequence)
+
+ def get_lagged_subsequences(
+ self, sequence: torch.Tensor, subsequences_length: int, shift: int = 0
+ ) -> torch.Tensor:
+ """
+ Returns lagged subsequences of a given sequence. Returns a tensor of shape (batch_size, subsequences_length,
+ feature_size, indices_length), containing lagged subsequences. Specifically, lagged[i, j, :, k] = sequence[i,
+ -indices[k]-subsequences_length+j, :].
+
+ Args:
+ sequence (`torch.Tensor` or shape `(batch_size, context_length,
+ feature_size)`): The sequence from which lagged subsequences should be extracted.
+ subsequences_length (`int`):
+ Length of the subsequences to be extracted.
+ shift (`int`, *optional* defaults to 0):
+ Shift the lags by this amount back in the time index.
+ """
+
+ # calculates the indices of the lags by subtracting the shift value from the given lags_sequence
+ indices = [lag - shift for lag in self.config.lags_sequence]
+
+ # checks if the maximum lag plus the length of the subsequences exceeds the length of the input sequence
+ sequence_length = sequence.shape[1]
+ if max(indices) + subsequences_length > sequence_length:
+ raise ValueError(
+ f"lags cannot go further than history length, found lag {max(indices)} "
+ f"while history length is only {sequence_length}"
+ )
+
+ # extracts the lagged subsequences from the input sequence using the calculated indices
+ lagged_values = []
+ for lag_index in indices:
+ begin_index = -lag_index - subsequences_length
+ end_index = -lag_index if lag_index > 0 else None
+ lagged_values.append(sequence[:, begin_index:end_index, ...])
+
+ # return as stacked tensor in the feature dimension
+ return torch.stack(lagged_values, dim=-1)
+
+ def create_network_inputs(
+ self,
+ past_values: torch.Tensor,
+ past_time_features: torch.Tensor,
+ static_categorical_features: torch.Tensor | None = None,
+ static_real_features: torch.Tensor | None = None,
+ past_observed_mask: torch.Tensor | None = None,
+ future_values: torch.Tensor | None = None,
+ future_time_features: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Creates the inputs for the network given the past and future values, time features, and static features.
+
+ Args:
+ past_values (`torch.Tensor`):
+ A tensor of shape `(batch_size, past_length, input_size)` containing the past values.
+ past_time_features (`torch.Tensor`):
+ A tensor of shape `(batch_size, past_length, num_features)` containing the past time features.
+ static_categorical_features (`Optional[torch.Tensor]`):
+ An optional tensor of shape `(batch_size, num_categorical_features)` containing the static categorical
+ features.
+ static_real_features (`Optional[torch.Tensor]`):
+ An optional tensor of shape `(batch_size, num_real_features)` containing the static real features.
+ past_observed_mask (`Optional[torch.Tensor]`):
+ An optional tensor of shape `(batch_size, past_length, input_size)` containing the mask of observed
+ values in the past.
+ future_values (`Optional[torch.Tensor]`):
+ An optional tensor of shape `(batch_size, future_length, input_size)` containing the future values.
+
+ Returns:
+ A tuple containing the following tensors:
+ - reshaped_lagged_sequence (`torch.Tensor`): A tensor of shape `(batch_size, sequence_length, num_lags *
+ input_size)` containing the lagged subsequences of the inputs.
+ - features (`torch.Tensor`): A tensor of shape `(batch_size, sequence_length, num_features)` containing the
+ concatenated static and time features.
+ - loc (`torch.Tensor`): A tensor of shape `(batch_size, input_size)` containing the mean of the input
+ values.
+ - scale (`torch.Tensor`): A tensor of shape `(batch_size, input_size)` containing the std of the input
+ values.
+ - static_feat (`torch.Tensor`): A tensor of shape `(batch_size, num_static_features)` containing the
+ concatenated static features.
+ """
+ # time feature
+ time_feat = (
+ torch.cat(
+ (
+ past_time_features[:, self._past_length - self.config.context_length :, ...],
+ future_time_features,
+ ),
+ dim=1,
+ )
+ if future_values is not None
+ else past_time_features[:, self._past_length - self.config.context_length :, ...]
+ )
+
+ # target
+ if past_observed_mask is None:
+ past_observed_mask = torch.ones_like(past_values)
+
+ context = past_values[:, -self.config.context_length :]
+ observed_context = past_observed_mask[:, -self.config.context_length :]
+ _, loc, scale = self.scaler(context, observed_context)
+
+ inputs = (
+ (torch.cat((past_values, future_values), dim=1) - loc) / scale
+ if future_values is not None
+ else (past_values - loc) / scale
+ )
+
+ # static features
+ log_abs_loc = loc.abs().log1p() if self.config.input_size == 1 else loc.squeeze(1).abs().log1p()
+ log_scale = scale.log() if self.config.input_size == 1 else scale.squeeze(1).log()
+ static_feat = torch.cat((log_abs_loc, log_scale), dim=1)
+
+ if static_real_features is not None:
+ static_feat = torch.cat((static_real_features, static_feat), dim=1)
+ if static_categorical_features is not None:
+ embedded_cat = self.embedder(static_categorical_features)
+ static_feat = torch.cat((embedded_cat, static_feat), dim=1)
+ expanded_static_feat = static_feat.unsqueeze(1).expand(-1, time_feat.shape[1], -1)
+
+ # all features
+ features = torch.cat((expanded_static_feat, time_feat), dim=-1)
+
+ # lagged features
+ subsequences_length = (
+ self.config.context_length + self.config.prediction_length
+ if future_values is not None
+ else self.config.context_length
+ )
+ lagged_sequence = self.get_lagged_subsequences(sequence=inputs, subsequences_length=subsequences_length)
+ lags_shape = lagged_sequence.shape
+ reshaped_lagged_sequence = lagged_sequence.reshape(lags_shape[0], lags_shape[1], -1)
+
+ if reshaped_lagged_sequence.shape[1] != time_feat.shape[1]:
+ raise ValueError(
+ f"input length {reshaped_lagged_sequence.shape[1]} and time feature lengths {time_feat.shape[1]} does not match"
+ )
+ return reshaped_lagged_sequence, features, loc, scale, static_feat
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ past_values: torch.Tensor,
+ past_time_features: torch.Tensor,
+ past_observed_mask: torch.Tensor,
+ static_categorical_features: torch.Tensor | None = None,
+ static_real_features: torch.Tensor | None = None,
+ future_values: torch.Tensor | None = None,
+ future_time_features: torch.Tensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: list[torch.FloatTensor] | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> AutoformerModelOutput:
+ r"""
+ past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
+ Past values of the time series, that serve as context in order to predict the future. These values may
+ contain lags, i.e. additional values from the past which are added in order to serve as "extra context".
+ The `past_values` is what the Transformer encoder gets as input (with optional additional features, such as
+ `static_categorical_features`, `static_real_features`, `past_time_features`).
+
+ The sequence length here is equal to `context_length` + `max(config.lags_sequence)`.
+
+ Missing values need to be replaced with zeros.
+ past_time_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_features)`, *optional*):
+ Optional time features, which the model internally will add to `past_values`. These could be things like
+ "month of year", "day of the month", etc. encoded as vectors (for instance as Fourier features). These
+ could also be so-called "age" features, which basically help the model know "at which point in life" a
+ time-series is. Age features have small values for distant past time steps and increase monotonically the
+ more we approach the current time step.
+
+ These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, where
+ the position encodings are learned from scratch internally as parameters of the model, the Time Series
+ Transformer requires to provide additional time features.
+
+ The Autoformer only learns additional embeddings for `static_categorical_features`.
+ past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in
+ `[0, 1]`:
+
+ - 1 for values that are **observed**,
+ - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+ static_categorical_features (`torch.LongTensor` of shape `(batch_size, number of static categorical features)`, *optional*):
+ Optional static categorical features for which the model will learn an embedding, which it will add to the
+ values of the time series.
+
+ Static categorical features are features which have the same value for all time steps (static over time).
+
+ A typical example of a static categorical feature is a time series ID.
+ static_real_features (`torch.FloatTensor` of shape `(batch_size, number of static real features)`, *optional*):
+ Optional static real features which the model will add to the values of the time series.
+
+ Static real features are features which have the same value for all time steps (static over time).
+
+ A typical example of a static real feature is promotion information.
+ future_values (`torch.FloatTensor` of shape `(batch_size, prediction_length)`):
+ Future values of the time series, that serve as labels for the model. The `future_values` is what the
+ Transformer needs to learn to output, given the `past_values`.
+
+ See the demo notebook and code snippets for details.
+
+ Missing values need to be replaced with zeros.
+ future_time_features (`torch.FloatTensor` of shape `(batch_size, prediction_length, num_features)`, *optional*):
+ Optional time features, which the model internally will add to `future_values`. These could be things like
+ "month of year", "day of the month", etc. encoded as vectors (for instance as Fourier features). These
+ could also be so-called "age" features, which basically help the model know "at which point in life" a
+ time-series is. Age features have small values for distant past time steps and increase monotonically the
+ more we approach the current time step.
+
+ These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, where
+ the position encodings are learned from scratch internally as parameters of the model, the Time Series
+ Transformer requires to provide additional features.
+
+ The Autoformer only learns additional embeddings for `static_categorical_features`.
+ encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*):
+ Tuple consists of `last_hidden_state`, `hidden_states` (*optional*) and `attentions` (*optional*)
+ `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)` (*optional*) is a sequence of
+ hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder.
+
+ Examples:
+
+ ```python
+ >>> from huggingface_hub import hf_hub_download
+ >>> import torch
+ >>> from transformers import AutoformerModel
+
+ >>> file = hf_hub_download(
+ ... repo_id="hf-internal-testing/tourism-monthly-batch", filename="train-batch.pt", repo_type="dataset"
+ ... )
+ >>> batch = torch.load(file)
+
+ >>> model = AutoformerModel.from_pretrained("huggingface/autoformer-tourism-monthly")
+
+ >>> # during training, one provides both past and future values
+ >>> # as well as possible additional features
+ >>> outputs = model(
+ ... past_values=batch["past_values"],
+ ... past_time_features=batch["past_time_features"],
+ ... past_observed_mask=batch["past_observed_mask"],
+ ... static_categorical_features=batch["static_categorical_features"],
+ ... future_values=batch["future_values"],
+ ... future_time_features=batch["future_time_features"],
+ ... )
+
+ >>> last_hidden_state = outputs.last_hidden_state
+ ```"""
+ transformer_inputs, temporal_features, loc, scale, static_feat = self.create_network_inputs(
+ past_values=past_values,
+ past_time_features=past_time_features,
+ past_observed_mask=past_observed_mask,
+ static_categorical_features=static_categorical_features,
+ static_real_features=static_real_features,
+ future_values=future_values,
+ future_time_features=future_time_features,
+ )
+
+ if encoder_outputs is None:
+ enc_input = torch.cat(
+ (
+ transformer_inputs[:, : self.config.context_length, ...],
+ temporal_features[:, : self.config.context_length, ...],
+ ),
+ dim=-1,
+ )
+ encoder_outputs: BaseModelOutput = self.encoder(
+ inputs_embeds=enc_input,
+ **kwargs,
+ )
+ elif not isinstance(encoder_outputs, BaseModelOutput):
+ encoder_outputs = BaseModelOutput(
+ last_hidden_state=encoder_outputs[0],
+ hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
+ attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
+ )
+
+ if future_values is not None:
+ # Decoder inputs
+ # seasonality and trend from context length
+ seasonal_input, trend_input = self.decomposition_layer(
+ transformer_inputs[:, : self.config.context_length, ...]
+ )
+ mean = (
+ torch.mean(transformer_inputs[:, : self.config.context_length, ...], dim=1)
+ .unsqueeze(1)
+ .repeat(1, self.config.prediction_length, 1)
+ )
+ zeros = torch.zeros(
+ [transformer_inputs.shape[0], self.config.prediction_length, transformer_inputs.shape[2]],
+ device=enc_input.device,
+ )
+
+ decoder_input = torch.cat(
+ (
+ torch.cat((seasonal_input[:, -self.config.label_length :, ...], zeros), dim=1),
+ temporal_features[:, self.config.context_length - self.config.label_length :, ...],
+ ),
+ dim=-1,
+ )
+ trend_init = torch.cat(
+ (
+ torch.cat((trend_input[:, -self.config.label_length :, ...], mean), dim=1),
+ temporal_features[:, self.config.context_length - self.config.label_length :, ...],
+ ),
+ dim=-1,
+ )
+
+ decoder_outputs: AutoFormerDecoderOutput = self.decoder(
+ trend=trend_init,
+ inputs_embeds=decoder_input,
+ attention_mask=decoder_attention_mask,
+ encoder_hidden_states=encoder_outputs.last_hidden_state,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+ else:
+ decoder_outputs = AutoFormerDecoderOutput()
+
+ return AutoformerModelOutput(
+ last_hidden_state=decoder_outputs.last_hidden_state,
+ trend=decoder_outputs.trend,
+ past_key_values=decoder_outputs.past_key_values,
+ decoder_hidden_states=decoder_outputs.hidden_states,
+ decoder_attentions=decoder_outputs.attentions,
+ cross_attentions=decoder_outputs.cross_attentions,
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+ encoder_hidden_states=encoder_outputs.hidden_states,
+ encoder_attentions=encoder_outputs.attentions,
+ loc=loc,
+ scale=scale,
+ static_features=static_feat,
+ )
+
+
+@auto_docstring
+class AutoformerForPrediction(AutoformerPreTrainedModel):
+ def __init__(self, config: AutoformerConfig):
+ super().__init__(config)
+ self.model = AutoformerModel(config)
+ if config.distribution_output == "student_t":
+ self.distribution_output = StudentTOutput(dim=config.input_size)
+ elif config.distribution_output == "normal":
+ self.distribution_output = NormalOutput(dim=config.input_size)
+ elif config.distribution_output == "negative_binomial":
+ self.distribution_output = NegativeBinomialOutput(dim=config.input_size)
+ else:
+ raise ValueError(f"Unknown distribution output {config.distribution_output}")
+
+ self.parameter_projection = self.distribution_output.get_parameter_projection(self.model.config.feature_size)
+ self.target_shape = self.distribution_output.event_shape
+
+ if config.loss == "nll":
+ self.loss = nll
+ else:
+ raise ValueError(f"Unknown loss function {config.loss}")
+
+ # Initialize weights of distribution_output and apply final processing
+ self.post_init()
+
+ def output_params(self, decoder_output):
+ return self.parameter_projection(decoder_output[:, -self.config.prediction_length :, :])
+
+ @torch.jit.ignore
+ def output_distribution(self, params, loc=None, scale=None, trailing_n=None) -> torch.distributions.Distribution:
+ sliced_params = params
+ if trailing_n is not None:
+ sliced_params = [p[:, -trailing_n:] for p in params]
+ return self.distribution_output.distribution(sliced_params, loc=loc, scale=scale)
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ past_values: torch.Tensor,
+ past_time_features: torch.Tensor,
+ past_observed_mask: torch.Tensor,
+ static_categorical_features: torch.Tensor | None = None,
+ static_real_features: torch.Tensor | None = None,
+ future_values: torch.Tensor | None = None,
+ future_time_features: torch.Tensor | None = None,
+ future_observed_mask: torch.Tensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: list[torch.FloatTensor] | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Seq2SeqTSPredictionOutput:
+ r"""
+ past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
+ Past values of the time series, that serve as context in order to predict the future. These values may
+ contain lags, i.e. additional values from the past which are added in order to serve as "extra context".
+ The `past_values` is what the Transformer encoder gets as input (with optional additional features, such as
+ `static_categorical_features`, `static_real_features`, `past_time_features`).
+
+ The sequence length here is equal to `context_length` + `max(config.lags_sequence)`.
+
+ Missing values need to be replaced with zeros.
+ past_time_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_features)`, *optional*):
+ Optional time features, which the model internally will add to `past_values`. These could be things like
+ "month of year", "day of the month", etc. encoded as vectors (for instance as Fourier features). These
+ could also be so-called "age" features, which basically help the model know "at which point in life" a
+ time-series is. Age features have small values for distant past time steps and increase monotonically the
+ more we approach the current time step.
+
+ These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, where
+ the position encodings are learned from scratch internally as parameters of the model, the Time Series
+ Transformer requires to provide additional time features.
+
+ The Autoformer only learns additional embeddings for `static_categorical_features`.
+ past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected in
+ `[0, 1]`:
+
+ - 1 for values that are **observed**,
+ - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+ static_categorical_features (`torch.LongTensor` of shape `(batch_size, number of static categorical features)`, *optional*):
+ Optional static categorical features for which the model will learn an embedding, which it will add to the
+ values of the time series.
+
+ Static categorical features are features which have the same value for all time steps (static over time).
+
+ A typical example of a static categorical feature is a time series ID.
+ static_real_features (`torch.FloatTensor` of shape `(batch_size, number of static real features)`, *optional*):
+ Optional static real features which the model will add to the values of the time series.
+
+ Static real features are features which have the same value for all time steps (static over time).
+
+ A typical example of a static real feature is promotion information.
+ future_values (`torch.FloatTensor` of shape `(batch_size, prediction_length)`):
+ Future values of the time series, that serve as labels for the model. The `future_values` is what the
+ Transformer needs to learn to output, given the `past_values`.
+
+ See the demo notebook and code snippets for details.
+
+ Missing values need to be replaced with zeros.
+ future_time_features (`torch.FloatTensor` of shape `(batch_size, prediction_length, num_features)`, *optional*):
+ Optional time features, which the model internally will add to `future_values`. These could be things like
+ "month of year", "day of the month", etc. encoded as vectors (for instance as Fourier features). These
+ could also be so-called "age" features, which basically help the model know "at which point in life" a
+ time-series is. Age features have small values for distant past time steps and increase monotonically the
+ more we approach the current time step.
+
+ These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT, where
+ the position encodings are learned from scratch internally as parameters of the model, the Time Series
+ Transformer requires to provide additional features.
+
+ The Autoformer only learns additional embeddings for `static_categorical_features`.
+ future_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`, *optional*):
+ Boolean mask to indicate which `future_values` were observed and which were missing. Mask values selected
+ in `[0, 1]`:
+
+ - 1 for values that are **observed**,
+ - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+
+ This mask is used to filter out missing values for the final loss calculation.
+ encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*):
+ Tuple consists of `last_hidden_state`, `hidden_states` (*optional*) and `attentions` (*optional*)
+ `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)` (*optional*) is a sequence of
+ hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder.
+
+ Examples:
+
+ ```python
+ >>> from huggingface_hub import hf_hub_download
+ >>> import torch
+ >>> from transformers import AutoformerForPrediction
+
+ >>> file = hf_hub_download(
+ ... repo_id="hf-internal-testing/tourism-monthly-batch", filename="train-batch.pt", repo_type="dataset"
+ ... )
+ >>> batch = torch.load(file)
+
+ >>> model = AutoformerForPrediction.from_pretrained("huggingface/autoformer-tourism-monthly")
+
+ >>> # during training, one provides both past and future values
+ >>> # as well as possible additional features
+ >>> outputs = model(
+ ... past_values=batch["past_values"],
+ ... past_time_features=batch["past_time_features"],
+ ... past_observed_mask=batch["past_observed_mask"],
+ ... static_categorical_features=batch["static_categorical_features"],
+ ... future_values=batch["future_values"],
+ ... future_time_features=batch["future_time_features"],
+ ... )
+
+ >>> loss = outputs.loss
+ >>> loss.backward()
+
+ >>> # during inference, one only provides past values
+ >>> # as well as possible additional features
+ >>> # the model autoregressively generates future values
+ >>> outputs = model.generate(
+ ... past_values=batch["past_values"],
+ ... past_time_features=batch["past_time_features"],
+ ... past_observed_mask=batch["past_observed_mask"],
+ ... static_categorical_features=batch["static_categorical_features"],
+ ... future_time_features=batch["future_time_features"],
+ ... )
+
+ >>> mean_prediction = outputs.sequences.mean(dim=1)
+ ```
+
+
+
+ The AutoformerForPrediction can also use static_real_features. To do so, set num_static_real_features in
+ AutoformerConfig based on number of such features in the dataset (in case of tourism_monthly dataset it
+ is equal to 1), initialize the model and call as shown below:
+
+ ```
+ >>> from huggingface_hub import hf_hub_download
+ >>> import torch
+ >>> from transformers import AutoformerConfig, AutoformerForPrediction
+
+ >>> file = hf_hub_download(
+ ... repo_id="hf-internal-testing/tourism-monthly-batch", filename="train-batch.pt", repo_type="dataset"
+ ... )
+ >>> batch = torch.load(file)
+
+ >>> # check number of static real features
+ >>> num_static_real_features = batch["static_real_features"].shape[-1]
+
+ >>> # load configuration of pretrained model and override num_static_real_features
+ >>> configuration = AutoformerConfig.from_pretrained(
+ ... "huggingface/autoformer-tourism-monthly",
+ ... num_static_real_features=num_static_real_features,
+ ... )
+ >>> # we also need to update feature_size as it is not recalculated
+ >>> configuration.feature_size += num_static_real_features
+
+ >>> model = AutoformerForPrediction(configuration)
+
+ >>> outputs = model(
+ ... past_values=batch["past_values"],
+ ... past_time_features=batch["past_time_features"],
+ ... past_observed_mask=batch["past_observed_mask"],
+ ... static_categorical_features=batch["static_categorical_features"],
+ ... static_real_features=batch["static_real_features"],
+ ... future_values=batch["future_values"],
+ ... future_time_features=batch["future_time_features"],
+ ... )
+ ```
+
+
+ """
+
+ if future_values is not None:
+ use_cache = False
+
+ outputs: AutoformerModelOutput = self.model(
+ past_values=past_values,
+ past_time_features=past_time_features,
+ past_observed_mask=past_observed_mask,
+ static_categorical_features=static_categorical_features,
+ static_real_features=static_real_features,
+ future_values=future_values,
+ future_time_features=future_time_features,
+ decoder_attention_mask=decoder_attention_mask,
+ encoder_outputs=encoder_outputs,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ prediction_loss = None
+ params = None
+ if future_values is not None:
+ params = self.output_params(outputs.last_hidden_state + outputs.trend)
+ distribution = self.output_distribution(params, loc=outputs.loc, scale=outputs.scale)
+
+ loss = self.loss(distribution, future_values)
+
+ if future_observed_mask is None:
+ future_observed_mask = torch.ones_like(future_values)
+
+ if len(self.target_shape) == 0:
+ loss_weights = future_observed_mask
+ else:
+ loss_weights, _ = future_observed_mask.min(dim=-1, keepdim=False)
+
+ prediction_loss = weighted_average(loss, weights=loss_weights)
+
+ return Seq2SeqTSPredictionOutput(
+ loss=prediction_loss,
+ params=params,
+ past_key_values=outputs.past_key_values,
+ decoder_hidden_states=outputs.decoder_hidden_states,
+ decoder_attentions=outputs.decoder_attentions,
+ cross_attentions=outputs.cross_attentions,
+ encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+ encoder_hidden_states=outputs.encoder_hidden_states,
+ encoder_attentions=outputs.encoder_attentions,
+ loc=outputs.loc,
+ scale=outputs.scale,
+ static_features=outputs.static_features,
+ )
+
+ @torch.no_grad()
+ def generate(
+ self,
+ past_values: torch.Tensor,
+ past_time_features: torch.Tensor,
+ future_time_features: torch.Tensor,
+ past_observed_mask: torch.Tensor | None = None,
+ static_categorical_features: torch.Tensor | None = None,
+ static_real_features: torch.Tensor | None = None,
+ ) -> SampleTSPredictionOutput:
+ r"""
+ Greedily generate sequences of sample predictions from a model with a probability distribution head.
+
+ Parameters:
+ past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`):
+ Past values of the time series, that serve as context in order to predict the future. The sequence size
+ of this tensor must be larger than the `context_length` of the model, since the model will use the
+ larger size to construct lag features, i.e. additional values from the past which are added in order to
+ serve as "extra context".
+
+ The `sequence_length` here is equal to `config.context_length` + `max(config.lags_sequence)`, which if
+ no `lags_sequence` is configured, is equal to `config.context_length` + 7 (as by default, the largest
+ look-back index in `config.lags_sequence` is 7). The property `_past_length` returns the actual length
+ of the past.
+
+ The `past_values` is what the Transformer encoder gets as input (with optional additional features,
+ such as `static_categorical_features`, `static_real_features`, `past_time_features` and lags).
+
+ Optionally, missing values need to be replaced with zeros and indicated via the `past_observed_mask`.
+
+ For multivariate time series, the `input_size` > 1 dimension is required and corresponds to the number
+ of variates in the time series per time step.
+ past_time_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_features)`):
+ Required time features, which the model internally will add to `past_values`. These could be things
+ like "month of year", "day of the month", etc. encoded as vectors (for instance as Fourier features).
+ These could also be so-called "age" features, which basically help the model know "at which point in
+ life" a time-series is. Age features have small values for distant past time steps and increase
+ monotonically the more we approach the current time step. Holiday features are also a good example of
+ time features.
+
+ These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT,
+ where the position encodings are learned from scratch internally as parameters of the model, the Time
+ Series Transformer requires to provide additional time features. The Time Series Transformer only
+ learns additional embeddings for `static_categorical_features`.
+
+ Additional dynamic real covariates can be concatenated to this tensor, with the caveat that these
+ features must but known at prediction time.
+
+ The `num_features` here is equal to `config.`num_time_features` + `config.num_dynamic_real_features`.
+ future_time_features (`torch.FloatTensor` of shape `(batch_size, prediction_length, num_features)`):
+ Required time features for the prediction window, which the model internally will add to sampled
+ predictions. These could be things like "month of year", "day of the month", etc. encoded as vectors
+ (for instance as Fourier features). These could also be so-called "age" features, which basically help
+ the model know "at which point in life" a time-series is. Age features have small values for distant
+ past time steps and increase monotonically the more we approach the current time step. Holiday features
+ are also a good example of time features.
+
+ These features serve as the "positional encodings" of the inputs. So contrary to a model like BERT,
+ where the position encodings are learned from scratch internally as parameters of the model, the Time
+ Series Transformer requires to provide additional time features. The Time Series Transformer only
+ learns additional embeddings for `static_categorical_features`.
+
+ Additional dynamic real covariates can be concatenated to this tensor, with the caveat that these
+ features must but known at prediction time.
+
+ The `num_features` here is equal to `config.`num_time_features` + `config.num_dynamic_real_features`.
+ past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length)` or `(batch_size, sequence_length, input_size)`, *optional*):
+ Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected
+ in `[0, 1]`:
+
+ - 1 for values that are **observed**,
+ - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+
+ static_categorical_features (`torch.LongTensor` of shape `(batch_size, number of static categorical features)`, *optional*):
+ Optional static categorical features for which the model will learn an embedding, which it will add to
+ the values of the time series.
+
+ Static categorical features are features which have the same value for all time steps (static over
+ time).
+
+ A typical example of a static categorical feature is a time series ID.
+ static_real_features (`torch.FloatTensor` of shape `(batch_size, number of static real features)`, *optional*):
+ Optional static real features which the model will add to the values of the time series.
+
+ Static real features are features which have the same value for all time steps (static over time).
+
+ A typical example of a static real feature is promotion information.
+ Return:
+ [`SampleTSPredictionOutput`] where the outputs `sequences` tensor will have shape `(batch_size, number of
+ samples, prediction_length)` or `(batch_size, number of samples, prediction_length, input_size)` for
+ multivariate predictions.
+ """
+ outputs = self(
+ static_categorical_features=static_categorical_features,
+ static_real_features=static_real_features,
+ past_time_features=past_time_features,
+ past_values=past_values,
+ past_observed_mask=past_observed_mask,
+ future_time_features=None,
+ future_values=None,
+ use_cache=False,
+ )
+
+ decoder = self.model.get_decoder()
+ enc_last_hidden = outputs.encoder_last_hidden_state
+ loc = outputs.loc
+ scale = outputs.scale
+ static_feat = outputs.static_features
+
+ num_parallel_samples = self.config.num_parallel_samples
+ repeated_loc = loc.repeat_interleave(repeats=num_parallel_samples, dim=0)
+ repeated_scale = scale.repeat_interleave(repeats=num_parallel_samples, dim=0)
+
+ repeated_past_values = (
+ past_values.repeat_interleave(repeats=num_parallel_samples, dim=0) - repeated_loc
+ ) / repeated_scale
+
+ time_features = torch.cat((past_time_features, future_time_features), dim=1)
+
+ expanded_static_feat = static_feat.unsqueeze(1).expand(-1, time_features.shape[1], -1)
+ features = torch.cat((expanded_static_feat, time_features), dim=-1)
+ repeated_features = features.repeat_interleave(repeats=num_parallel_samples, dim=0)
+
+ repeated_enc_last_hidden = enc_last_hidden.repeat_interleave(repeats=num_parallel_samples, dim=0)
+
+ lagged_sequence = self.model.get_lagged_subsequences(
+ sequence=repeated_past_values, subsequences_length=self.config.context_length
+ )
+ lags_shape = lagged_sequence.shape
+ reshaped_lagged_sequence = lagged_sequence.reshape(lags_shape[0], lags_shape[1], -1)
+ seasonal_input, trend_input = self.model.decomposition_layer(reshaped_lagged_sequence)
+
+ mean = torch.mean(reshaped_lagged_sequence, dim=1).unsqueeze(1).repeat(1, self.config.prediction_length, 1)
+ zeros = torch.zeros(
+ [reshaped_lagged_sequence.shape[0], self.config.prediction_length, reshaped_lagged_sequence.shape[2]],
+ device=reshaped_lagged_sequence.device,
+ )
+
+ decoder_input = torch.cat(
+ (
+ torch.cat((seasonal_input[:, -self.config.label_length :, ...], zeros), dim=1),
+ repeated_features[:, -self.config.prediction_length - self.config.label_length :, ...],
+ ),
+ dim=-1,
+ )
+ trend_init = torch.cat(
+ (
+ torch.cat((trend_input[:, -self.config.label_length :, ...], mean), dim=1),
+ repeated_features[:, -self.config.prediction_length - self.config.label_length :, ...],
+ ),
+ dim=-1,
+ )
+ decoder_outputs = decoder(
+ trend=trend_init, inputs_embeds=decoder_input, encoder_hidden_states=repeated_enc_last_hidden
+ )
+ decoder_last_hidden = decoder_outputs.last_hidden_state
+ trend = decoder_outputs.trend
+ params = self.output_params(decoder_last_hidden + trend)
+ distr = self.output_distribution(params, loc=repeated_loc, scale=repeated_scale)
+ future_samples = distr.sample()
+
+ return SampleTSPredictionOutput(
+ sequences=future_samples.reshape(
+ (-1, num_parallel_samples, self.config.prediction_length) + self.target_shape,
+ )
+ )
+
+
+__all__ = ["AutoformerForPrediction", "AutoformerModel", "AutoformerPreTrainedModel"]
diff --git a/third_party/transformers/src/transformers/models/bark/__init__.py b/third_party/transformers/src/transformers/models/bark/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c5296fc47fc423c300cf6c43bccb5daafd6d134a
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bark/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_bark import *
+ from .modeling_bark import *
+ from .processing_bark import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/bark/configuration_bark.py b/third_party/transformers/src/transformers/models/bark/configuration_bark.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae24fa2e74778eaee6be81a6ca1fc48f0998f295
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bark/configuration_bark.py
@@ -0,0 +1,275 @@
+# Copyright 2023 The Suno AI Authors and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""BARK model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="suno/bark")
+@strict
+class BarkSubModelConfig(PreTrainedConfig):
+ r"""
+ block_size (`int`, *optional*, defaults to 1024):
+ The maximum sequence length that this model might ever be used with. Typically set this to something large
+ just in case (e.g., 512 or 1024 or 2048).
+ input_vocab_size (`int`, *optional*, defaults to 10_048):
+ Vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`{model}`]. Defaults to 10_048 but should be carefully thought with
+ regards to the chosen sub-model.
+ output_vocab_size (`int`, *optional*, defaults to 10_048):
+ Output vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented
+ by the: `output_ids` when passing forward a [`{model}`]. Defaults to 10_048 but should be carefully thought
+ with regards to the chosen sub-model.
+ bias (`bool`, *optional*, defaults to `True`):
+ Whether or not to use bias in the linear layers and layer norm layers.
+ """
+
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ attribute_map = {
+ "num_attention_heads": "num_heads",
+ "num_hidden_layers": "num_layers",
+ "vocab_size": "input_vocab_size",
+ "window_size": "block_size",
+ }
+
+ block_size: int = 1024
+ input_vocab_size: int = 10_048
+ output_vocab_size: int = 10_048
+ num_layers: int = 12
+ num_heads: int = 12
+ hidden_size: int = 768
+ dropout: float | int = 0.0
+ bias: bool = True
+ initializer_range: float = 0.02
+ use_cache: bool = True
+
+
+@auto_docstring(checkpoint="suno/bark")
+@strict
+class BarkSemanticConfig(BarkSubModelConfig):
+ r"""
+ block_size (`int`, *optional*, defaults to 1024):
+ The maximum sequence length that this model might ever be used with. Typically set this to something large
+ just in case (e.g., 512 or 1024 or 2048).
+ input_vocab_size (`int`, *optional*, defaults to 10_048):
+ Vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`{model}`]. Defaults to 10_048 but should be carefully thought with
+ regards to the chosen sub-model.
+ output_vocab_size (`int`, *optional*, defaults to 10_048):
+ Output vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented
+ by the: `output_ids` when passing forward a [`{model}`]. Defaults to 10_048 but should be carefully thought
+ with regards to the chosen sub-model.
+ bias (`bool`, *optional*, defaults to `True`):
+ Whether or not to use bias in the linear layers and layer norm layers
+
+ Example:
+
+ ```python
+ >>> from transformers import BarkSemanticConfig, BarkSemanticModel
+
+ >>> # Initializing a Bark sub-module style configuration
+ >>> configuration = BarkSemanticConfig()
+
+ >>> # Initializing a model (with random weights) from the suno/bark style configuration
+ >>> model = BarkSemanticModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "semantic"
+ base_config_key = "semantic_config"
+
+
+@auto_docstring(checkpoint="suno/bark")
+@strict
+class BarkCoarseConfig(BarkSubModelConfig):
+ r"""
+ block_size (`int`, *optional*, defaults to 1024):
+ The maximum sequence length that this model might ever be used with. Typically set this to something large
+ just in case (e.g., 512 or 1024 or 2048).
+ input_vocab_size (`int`, *optional*, defaults to 10_048):
+ Vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`{model}`]. Defaults to 10_048 but should be carefully thought with
+ regards to the chosen sub-model.
+ output_vocab_size (`int`, *optional*, defaults to 10_048):
+ Output vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented
+ by the: `output_ids` when passing forward a [`{model}`]. Defaults to 10_048 but should be carefully thought
+ with regards to the chosen sub-model.
+ bias (`bool`, *optional*, defaults to `True`):
+ Whether or not to use bias in the linear layers and layer norm layers
+
+ Example:
+
+ ```python
+ >>> from transformers import BarkCoarseConfig, BarkCoarseModel
+
+ >>> # Initializing a Bark sub-module style configuration
+ >>> configuration = BarkCoarseConfig()
+
+ >>> # Initializing a model (with random weights) from the suno/bark style configuration
+ >>> model = BarkCoarseModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "coarse_acoustics"
+ base_config_key = "coarse_acoustics_config"
+
+
+@auto_docstring(checkpoint="suno/bark")
+@strict
+class BarkFineConfig(BarkSubModelConfig):
+ r"""
+ block_size (`int`, *optional*, defaults to 1024):
+ The maximum sequence length that this model might ever be used with. Typically set this to something large
+ just in case (e.g., 512 or 1024 or 2048).
+ input_vocab_size (`int`, *optional*, defaults to 10_048):
+ Vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented by the
+ `inputs_ids` passed when calling [`{model}`]. Defaults to 10_048 but should be carefully thought with
+ regards to the chosen sub-model.
+ output_vocab_size (`int`, *optional*, defaults to 10_048):
+ Output vocabulary size of a Bark sub-model. Defines the number of different tokens that can be represented
+ by the: `output_ids` when passing forward a [`{model}`]. Defaults to 10_048 but should be carefully thought
+ with regards to the chosen sub-model.
+ bias (`bool`, *optional*, defaults to `True`):
+ Whether or not to use bias in the linear layers and layer norm layers
+ n_codes_total (`int`, *optional*, defaults to 8):
+ The total number of audio codebooks predicted. Used in the fine acoustics sub-model.
+ n_codes_given (`int`, *optional*, defaults to 1):
+ The number of audio codebooks predicted in the coarse acoustics sub-model. Used in the acoustics
+ sub-models.
+
+ Example:
+
+ ```python
+ >>> from transformers import BarkFineConfig, BarkFineModel
+
+ >>> # Initializing a Bark sub-module style configuration
+ >>> configuration = BarkFineConfig()
+
+ >>> # Initializing a model (with random weights) from the suno/bark style configuration
+ >>> model = BarkFineModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "fine_acoustics"
+ base_config_key = "fine_acoustics_config"
+
+ tie_word_embeddings: bool = True
+ n_codes_total: int = 8
+ n_codes_given: int = 1
+
+
+@auto_docstring(checkpoint="suno/bark")
+@strict
+class BarkConfig(PreTrainedConfig):
+ r"""
+ semantic_config ([`BarkSemanticConfig`], *optional*):
+ Configuration of the underlying semantic sub-model.
+ coarse_acoustics_config ([`BarkCoarseConfig`], *optional*):
+ Configuration of the underlying coarse acoustics sub-model.
+ fine_acoustics_config ([`BarkFineConfig`], *optional*):
+ Configuration of the underlying fine acoustics sub-model.
+ codec_config ([`AutoConfig`], *optional*):
+ Configuration of the underlying codec sub-model.
+
+ Example:
+
+ ```python
+ >>> from transformers import (
+ ... BarkSemanticConfig,
+ ... BarkCoarseConfig,
+ ... BarkFineConfig,
+ ... BarkModel,
+ ... BarkConfig,
+ ... AutoConfig,
+ ... )
+
+ >>> # Initializing Bark sub-modules configurations.
+ >>> semantic_config = BarkSemanticConfig()
+ >>> coarse_acoustics_config = BarkCoarseConfig()
+ >>> fine_acoustics_config = BarkFineConfig()
+ >>> codec_config = AutoConfig.from_pretrained("facebook/encodec_24khz")
+
+
+ >>> # Initializing a Bark module style configuration
+ >>> configuration = BarkConfig(
+ ... semantic_config, coarse_acoustics_config, fine_acoustics_config, codec_config
+ ... )
+
+ >>> # Initializing a model (with random weights)
+ >>> model = BarkModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+ """
+
+ model_type = "bark"
+ sub_configs = {
+ "semantic_config": BarkSemanticConfig,
+ "coarse_acoustics_config": BarkCoarseConfig,
+ "fine_acoustics_config": BarkFineConfig,
+ "codec_config": AutoConfig,
+ }
+ semantic_config: dict | PreTrainedConfig | None = None
+ coarse_acoustics_config: dict | PreTrainedConfig | None = None
+ fine_acoustics_config: dict | PreTrainedConfig | None = None
+ codec_config: dict | PreTrainedConfig | None = None
+ initializer_range: float = 0.02
+
+ def __post_init__(self, **kwargs):
+ if self.semantic_config is None:
+ self.semantic_config = BarkSemanticConfig()
+ logger.info("`semantic_config` is `None`. Initializing the `BarkSemanticConfig` with default values.")
+ elif isinstance(self.semantic_config, dict):
+ self.semantic_config = BarkSemanticConfig(**self.semantic_config)
+
+ if self.coarse_acoustics_config is None:
+ self.coarse_acoustics_config = BarkCoarseConfig()
+ logger.info(
+ "`coarse_acoustics_config` is `None`. Initializing the `BarkCoarseConfig` with default values."
+ )
+ elif isinstance(self.coarse_acoustics_config, dict):
+ self.coarse_acoustics_config = BarkCoarseConfig(**self.coarse_acoustics_config)
+
+ if self.fine_acoustics_config is None:
+ self.fine_acoustics_config = BarkFineConfig()
+ logger.info("`fine_acoustics_config` is `None`. Initializing the `BarkFineConfig` with default values.")
+ elif isinstance(self.fine_acoustics_config, dict):
+ self.fine_acoustics_config = BarkFineConfig(**self.fine_acoustics_config)
+
+ if self.codec_config is None:
+ self.codec_config = CONFIG_MAPPING["encodec"]()
+ logger.info("`codec_config` is `None`. Initializing the `codec_config` with default values.")
+ elif isinstance(self.codec_config, dict):
+ codec_model_type = self.codec_config.get("model_type", "encodec")
+ self.codec_config = CONFIG_MAPPING[codec_model_type](**self.codec_config)
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["BarkCoarseConfig", "BarkConfig", "BarkFineConfig", "BarkSemanticConfig"]
diff --git a/third_party/transformers/src/transformers/models/bark/convert_suno_to_hf.py b/third_party/transformers/src/transformers/models/bark/convert_suno_to_hf.py
new file mode 100644
index 0000000000000000000000000000000000000000..d1c2f85a3c7b6f2191bf522100560f200b7ddf8f
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bark/convert_suno_to_hf.py
@@ -0,0 +1,261 @@
+"""Convert Bark checkpoint."""
+
+import argparse
+import os
+from pathlib import Path
+
+import torch
+from bark.generation import _load_model as _bark_load_model
+from huggingface_hub import hf_hub_download
+
+from transformers import EncodecConfig, EncodecModel, set_seed
+from transformers.models.bark.configuration_bark import (
+ BarkCoarseConfig,
+ BarkConfig,
+ BarkFineConfig,
+ BarkSemanticConfig,
+)
+from transformers.models.bark.generation_configuration_bark import (
+ BarkCoarseGenerationConfig,
+ BarkFineGenerationConfig,
+ BarkGenerationConfig,
+ BarkSemanticGenerationConfig,
+)
+from transformers.models.bark.modeling_bark import BarkCoarseModel, BarkFineModel, BarkModel, BarkSemanticModel
+from transformers.utils import logging
+
+
+logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+set_seed(770)
+
+
+new_layer_name_dict = {
+ "c_attn": "att_proj",
+ "c_proj": "out_proj",
+ "c_fc": "in_proj",
+ "transformer.": "",
+ "h.": "layers.",
+ "ln_1": "layernorm_1",
+ "ln_2": "layernorm_2",
+ "ln_f": "layernorm_final",
+ "wpe": "position_embeds_layer",
+ "wte": "input_embeds_layer",
+}
+
+
+REMOTE_MODEL_PATHS = {
+ "text_small": {
+ "repo_id": "suno/bark",
+ "file_name": "text.pt",
+ },
+ "coarse_small": {
+ "repo_id": "suno/bark",
+ "file_name": "coarse.pt",
+ },
+ "fine_small": {
+ "repo_id": "suno/bark",
+ "file_name": "fine.pt",
+ },
+ "text": {
+ "repo_id": "suno/bark",
+ "file_name": "text_2.pt",
+ },
+ "coarse": {
+ "repo_id": "suno/bark",
+ "file_name": "coarse_2.pt",
+ },
+ "fine": {
+ "repo_id": "suno/bark",
+ "file_name": "fine_2.pt",
+ },
+}
+
+CUR_PATH = os.path.dirname(os.path.abspath(__file__))
+default_cache_dir = os.path.join(os.path.expanduser("~"), ".cache")
+CACHE_DIR = os.path.join(os.getenv("XDG_CACHE_HOME", default_cache_dir), "suno", "bark_v0")
+
+
+def _get_ckpt_path(model_type, use_small=False):
+ key = model_type
+ if use_small:
+ key += "_small"
+ return os.path.join(CACHE_DIR, REMOTE_MODEL_PATHS[key]["file_name"])
+
+
+def _download(from_hf_path, file_name):
+ os.makedirs(CACHE_DIR, exist_ok=True)
+ hf_hub_download(repo_id=from_hf_path, filename=file_name, local_dir=CACHE_DIR)
+
+
+def _load_model(ckpt_path, device, use_small=False, model_type="text"):
+ if model_type == "text":
+ ModelClass = BarkSemanticModel
+ ConfigClass = BarkSemanticConfig
+ GenerationConfigClass = BarkSemanticGenerationConfig
+ elif model_type == "coarse":
+ ModelClass = BarkCoarseModel
+ ConfigClass = BarkCoarseConfig
+ GenerationConfigClass = BarkCoarseGenerationConfig
+ elif model_type == "fine":
+ ModelClass = BarkFineModel
+ ConfigClass = BarkFineConfig
+ GenerationConfigClass = BarkFineGenerationConfig
+ else:
+ raise NotImplementedError()
+ model_key = f"{model_type}_small" if use_small else model_type
+ model_info = REMOTE_MODEL_PATHS[model_key]
+ if not os.path.exists(ckpt_path):
+ logger.info(f"{model_type} model not found, downloading into `{CACHE_DIR}`.")
+ _download(model_info["repo_id"], model_info["file_name"])
+ checkpoint = torch.load(ckpt_path, map_location=device, weights_only=True)
+ # this is a hack
+ model_args = checkpoint["model_args"]
+ if "input_vocab_size" not in model_args:
+ model_args["input_vocab_size"] = model_args["vocab_size"]
+ model_args["output_vocab_size"] = model_args["vocab_size"]
+ del model_args["vocab_size"]
+
+ # convert Bark model arguments to HF Bark model arguments
+ model_args["num_heads"] = model_args.pop("n_head")
+ model_args["hidden_size"] = model_args.pop("n_embd")
+ model_args["num_layers"] = model_args.pop("n_layer")
+
+ model_config = ConfigClass(**checkpoint["model_args"])
+ model = ModelClass(config=model_config)
+ model_generation_config = GenerationConfigClass()
+
+ model.generation_config = model_generation_config
+ state_dict = checkpoint["model"]
+ # fixup checkpoint
+ unwanted_prefix = "_orig_mod."
+ for k in state_dict:
+ if k.startswith(unwanted_prefix):
+ # replace part of the key with corresponding layer name in HF implementation
+ new_k = k[len(unwanted_prefix) :]
+ for old_layer_name, new_layer_name in new_layer_name_dict.items():
+ new_k = new_k.replace(old_layer_name, new_layer_name)
+
+ state_dict[new_k] = state_dict.pop(k)
+
+ extra_keys = set(state_dict.keys()) - set(model.state_dict().keys())
+ extra_keys = {k for k in extra_keys if not k.endswith(".attn.bias")}
+ missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
+ missing_keys = {k for k in missing_keys if not k.endswith(".attn.bias")}
+ if len(extra_keys) != 0:
+ raise ValueError(f"extra keys found: {extra_keys}")
+ if len(missing_keys) != 0:
+ raise ValueError(f"missing keys: {missing_keys}")
+ model.load_state_dict(state_dict, strict=False)
+ n_params = model.num_parameters(exclude_embeddings=True)
+ val_loss = checkpoint["best_val_loss"].item()
+ logger.info(f"model loaded: {round(n_params / 1e6, 1)}M params, {round(val_loss, 3)} loss")
+ model.eval()
+ model.to(device)
+ del checkpoint, state_dict
+
+ return model
+
+
+def load_model(pytorch_dump_folder_path, use_small=False, model_type="text"):
+ if model_type not in ("text", "coarse", "fine"):
+ raise NotImplementedError()
+
+ device = "cpu" # do conversion on cpu
+
+ ckpt_path = _get_ckpt_path(model_type, use_small=use_small)
+ model = _load_model(ckpt_path, device, model_type=model_type, use_small=use_small)
+
+ # load bark initial model
+ bark_model = _bark_load_model(ckpt_path, "cpu", model_type=model_type, use_small=use_small)
+
+ if model_type == "text":
+ bark_model = bark_model["model"]
+
+ if model.num_parameters(exclude_embeddings=True) != bark_model.get_num_params():
+ raise ValueError("initial and new models don't have the same number of parameters")
+
+ # check if same output as the bark model
+ batch_size = 5
+ sequence_length = 10
+
+ if model_type in ["text", "coarse"]:
+ vec = torch.randint(256, (batch_size, sequence_length), dtype=torch.int)
+ output_old_model = bark_model(vec)[0]
+
+ output_new_model_total = model(vec)
+
+ # take last logits
+ output_new_model = output_new_model_total.logits[:, [-1], :]
+
+ else:
+ prediction_codebook_channel = 3
+ n_codes_total = 8
+ vec = torch.randint(256, (batch_size, sequence_length, n_codes_total), dtype=torch.int)
+
+ output_new_model_total = model(prediction_codebook_channel, vec)
+ output_old_model = bark_model(prediction_codebook_channel, vec)
+
+ output_new_model = output_new_model_total.logits
+
+ # output difference should come from the difference of self-attention implementation design
+ if output_new_model.shape != output_old_model.shape:
+ raise ValueError("initial and new outputs don't have the same shape")
+ if (output_new_model - output_old_model).abs().max().item() > 1e-3:
+ raise ValueError("initial and new outputs are not equal")
+
+ Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
+ model.save_pretrained(pytorch_dump_folder_path)
+
+
+def load_whole_bark_model(
+ semantic_path,
+ coarse_path,
+ fine_path,
+ append_text,
+ hub_path,
+ folder_path,
+):
+ pytorch_dump_folder_path = os.path.join(folder_path, append_text)
+
+ semanticConfig = BarkSemanticConfig.from_pretrained(os.path.join(semantic_path, "config.json"))
+ coarseAcousticConfig = BarkCoarseConfig.from_pretrained(os.path.join(coarse_path, "config.json"))
+ fineAcousticConfig = BarkFineConfig.from_pretrained(os.path.join(fine_path, "config.json"))
+ codecConfig = EncodecConfig.from_pretrained("facebook/encodec_24khz")
+
+ semantic = BarkSemanticModel.from_pretrained(semantic_path)
+ coarseAcoustic = BarkCoarseModel.from_pretrained(coarse_path)
+ fineAcoustic = BarkFineModel.from_pretrained(fine_path)
+ codec = EncodecModel.from_pretrained("facebook/encodec_24khz")
+
+ bark_config = BarkConfig(semanticConfig, coarseAcousticConfig, fineAcousticConfig, codecConfig)
+
+ bark_generation_config = BarkGenerationConfig(
+ semantic.generation_config, coarseAcoustic.generation_config, fineAcoustic.generation_config
+ )
+
+ bark = BarkModel(bark_config)
+
+ bark.semantic = semantic
+ bark.coarse_acoustics = coarseAcoustic
+ bark.fine_acoustics = fineAcoustic
+ bark.codec_model = codec
+
+ bark.generation_config = bark_generation_config
+
+ Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
+ bark.save_pretrained(pytorch_dump_folder_path, repo_id=hub_path, push_to_hub=True)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ # Required parameters
+
+ parser.add_argument("model_type", type=str, help="text, coarse or fine.")
+ parser.add_argument("pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
+ parser.add_argument("--is_small", action="store_true", help="convert the small version instead of the large.")
+
+ args = parser.parse_args()
+
+ load_model(args.pytorch_dump_folder_path, model_type=args.model_type, use_small=args.is_small)
diff --git a/third_party/transformers/src/transformers/models/bark/modeling_bark.py b/third_party/transformers/src/transformers/models/bark/modeling_bark.py
new file mode 100644
index 0000000000000000000000000000000000000000..e688affe2cdd7f6de684bee29ec7ecc29cf41b65
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bark/modeling_bark.py
@@ -0,0 +1,1521 @@
+# Copyright 2023 The Suno AI Authors and The HuggingFace Inc. team. All rights reserved.
+#
+# 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 BARK model."""
+
+import math
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from ... import initialization as init
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...generation.logits_process import (
+ AlternatingCodebooksLogitsProcessor,
+ BarkEosPrioritizerLogitsProcessor,
+ SuppressTokensLogitsProcessor,
+)
+from ...masking_utils import create_bidirectional_mask
+from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import CausalLMOutputWithPast, MaskedLMOutput
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
+ auto_docstring,
+ is_accelerate_available,
+ is_torch_accelerator_available,
+ logging,
+)
+from ...utils.deprecation import deprecate_kwarg
+from ..auto import AutoModel
+from .configuration_bark import (
+ BarkCoarseConfig,
+ BarkConfig,
+ BarkFineConfig,
+ BarkSemanticConfig,
+ BarkSubModelConfig,
+)
+from .generation_configuration_bark import (
+ BarkCoarseGenerationConfig,
+ BarkFineGenerationConfig,
+ BarkSemanticGenerationConfig,
+)
+
+
+if is_flash_attn_available():
+ from ...integrations.flash_attention import get_target_dtype
+ from ...modeling_flash_attention_utils import _flash_attention_forward
+
+
+logger = logging.get_logger(__name__)
+
+
+class BarkSelfAttention(nn.Module):
+ # adapted from GPTNeoSelfAttention and Bark code
+ # BarkSelfAttention can have two attention type, i.e full attention or causal attention
+
+ def __init__(self, config, is_causal=False, layer_idx=None):
+ super().__init__()
+
+ # regularization
+ self.dropout = config.dropout
+ self.attn_dropout = nn.Dropout(config.dropout)
+ self.resid_dropout = nn.Dropout(config.dropout)
+
+ self.embed_dim = config.hidden_size
+ self.num_heads = config.num_heads
+ self.head_dim = self.embed_dim // self.num_heads
+ self.config = config
+
+ if config.hidden_size % config.num_heads != 0:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+ f" {self.num_heads})."
+ )
+
+ # key, query, value projections for all heads, but in a batch
+ self.att_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=config.bias)
+ # output projection
+ self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.bias)
+
+ self.is_causal = is_causal
+ self.layer_idx = layer_idx
+ if is_causal:
+ block_size = config.block_size
+ bias = torch.tril(torch.ones((block_size, block_size), dtype=bool)).view(1, 1, block_size, block_size)
+ self.register_buffer("bias", bias)
+
+ # Copied from transformers.models.gpt_neo.modeling_gpt_neo.GPTNeoSelfAttention._split_heads
+ def _split_heads(self, tensor, num_heads, attn_head_size):
+ """
+ Splits hidden_size dim into attn_head_size and num_heads
+ """
+ new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
+ tensor = tensor.view(new_shape)
+ return tensor.permute(0, 2, 1, 3) # (batch, head, seq_length, head_features)
+
+ def _merge_heads(self, tensor, num_heads, attn_head_size):
+ """
+ Merges attn_head_size dim and num_attn_heads dim into hidden_size
+ """
+
+ # re-assemble all head outputs side by side
+ # (batch, num_heads, seq_len, attn_head_size) -> (batch, seq_len, num_heads*attn_head_size)
+ tensor = tensor.transpose(1, 2).contiguous()
+ tensor = tensor.view(tensor.size()[:-2] + (num_heads * attn_head_size,))
+
+ return tensor
+
+ def _attn(self, query, key, value, attention_mask=None):
+ # unlike GPTNeo's SelfAttention, divide by the square root of the dimension of the query and the key
+ attn_weights = torch.matmul(query, key.transpose(-1, -2)) * (1.0 / math.sqrt(self.head_dim))
+
+ if self.is_causal:
+ query_length, key_length = query.size(-2), key.size(-2)
+
+ # fill the upper left part of the attention weights with inf
+ attn_weights = attn_weights.masked_fill(
+ self.bias[:, :, key_length - query_length : key_length, :key_length] == 0,
+ torch.finfo(attn_weights.dtype).min,
+ )
+
+ if attention_mask is not None:
+ # Apply the attention mask
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ attn_weights = attn_weights.to(value.dtype)
+ attn_weights = self.attn_dropout(attn_weights)
+
+ # (batch, num_heads, seq_len, seq_len) x (batch, num_heads, seq_len, attn_head_size)
+ # -> (batch, num_heads, seq_len, attn_head_size)
+ attn_output = torch.matmul(attn_weights, value)
+
+ return attn_output, attn_weights
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ past_key_values=None,
+ use_cache=False,
+ output_attentions=False,
+ **kwargs,
+ ):
+ # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+ query, key, value = self.att_proj(hidden_states).split(self.embed_dim, dim=2)
+
+ query = self._split_heads(query, self.num_heads, self.head_dim)
+ key = self._split_heads(key, self.num_heads, self.head_dim)
+ value = self._split_heads(value, self.num_heads, self.head_dim)
+
+ if past_key_values is not None:
+ key, value = past_key_values.update(key, value, self.layer_idx)
+
+ attn_output, attn_weights = self._attn(query, key, value, attention_mask)
+
+ attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
+ attn_output = self.out_proj(attn_output)
+ attn_output = self.resid_dropout(attn_output)
+
+ return attn_output, attn_weights
+
+
+class BarkSelfFlashAttention2(BarkSelfAttention):
+ """
+ Bark flash attention module. This module inherits from `BarkSelfAttention` 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 __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+ # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+ self._flash_attn_uses_top_left_mask = flash_attn_supports_top_left_mask()
+
+ def _split_heads(self, tensor, num_heads, attn_head_size):
+ """
+ Splits hidden_size dim into attn_head_size and num_heads
+ """
+ new_shape = tensor.size()[:-1] + (num_heads, attn_head_size)
+ tensor = tensor.view(new_shape)
+ # Flash attention requires the input to have the shape
+ # batch_size x seq_length x head_dim x hidden_dim - (batch, seq_length, head, head_features)
+ return tensor
+
+ def _merge_heads(self, tensor, num_heads, attn_head_size):
+ """
+ Merges attn_head_size dim and num_attn_heads dim into hidden_size
+ """
+ # re-assemble all head outputs side by side
+ # (batch, seq_len, num_heads, attn_head_size) -> (batch, seq_len, num_heads*attn_head_size)
+ tensor = tensor.view(tensor.size()[:-2] + (num_heads * attn_head_size,))
+ return tensor
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ past_key_values=None,
+ use_cache=False,
+ output_attentions=False,
+ **kwargs,
+ ):
+ batch_size, query_len, _ = hidden_states.size()
+
+ # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+ query, key, value = self.att_proj(hidden_states).split(self.embed_dim, dim=2)
+
+ query = self._split_heads(query, self.num_heads, self.head_dim)
+ key = self._split_heads(key, self.num_heads, self.head_dim)
+ value = self._split_heads(value, self.num_heads, self.head_dim)
+
+ if past_key_values is not None:
+ key, value = past_key_values.update(key, value, self.layer_idx)
+
+ target_dtype = get_target_dtype(query, self) # if the query is in float32, this is the dtype to cast to for FA
+
+ attn_output = _flash_attention_forward(
+ query,
+ key,
+ value,
+ attention_mask,
+ query_len,
+ dropout=self.dropout if self.training else 0.0,
+ use_top_left_mask=self._flash_attn_uses_top_left_mask,
+ is_causal=self.is_causal,
+ target_dtype=target_dtype,
+ )
+
+ attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
+ attn_output = self.out_proj(attn_output)
+ attn_output = self.resid_dropout(attn_output)
+
+ return attn_output, None
+
+
+BARK_ATTENTION_CLASSES = {
+ "eager": BarkSelfAttention,
+ "flash_attention_2": BarkSelfFlashAttention2,
+}
+
+
+class BarkMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.in_proj = nn.Linear(config.hidden_size, 4 * config.hidden_size, bias=config.bias)
+ self.out_proj = nn.Linear(4 * config.hidden_size, config.hidden_size, bias=config.bias)
+ self.dropout = nn.Dropout(config.dropout)
+ self.gelu = nn.GELU()
+
+ def forward(self, hidden_states):
+ hidden_states = self.in_proj(hidden_states)
+ hidden_states = self.gelu(hidden_states)
+ hidden_states = self.out_proj(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ return hidden_states
+
+
+class BarkBlock(GradientCheckpointingLayer):
+ def __init__(self, config, is_causal=False, layer_idx=None):
+ super().__init__()
+
+ if is_causal:
+ # if causal, the layerNorm bias is optional to stick with Bark choice of leaving optional bias
+ # in AutoRegressive models (corresponding to the "Text" and the "Coarse" modules)
+ self.layernorm_1 = nn.LayerNorm(config.hidden_size, bias=config.bias)
+ self.layernorm_2 = nn.LayerNorm(config.hidden_size, bias=config.bias)
+ else:
+ self.layernorm_1 = nn.LayerNorm(config.hidden_size)
+ self.layernorm_2 = nn.LayerNorm(config.hidden_size)
+
+ self.attn = BARK_ATTENTION_CLASSES[config._attn_implementation](
+ config, is_causal=is_causal, layer_idx=layer_idx
+ )
+
+ self.mlp = BarkMLP(config)
+
+ def forward(
+ self,
+ hidden_states,
+ past_key_values=None,
+ attention_mask=None,
+ use_cache=False,
+ output_attentions=False,
+ **kwargs,
+ ):
+ intermediary_hidden_states = self.layernorm_1(hidden_states)
+
+ attn_outputs = self.attn(
+ intermediary_hidden_states,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ )
+
+ attn_output = attn_outputs[0] # output_attn: output, present_key_values, (attn_weights)
+ outputs = attn_outputs[1:]
+
+ intermediary_hidden_states = hidden_states + attn_output
+ intermediary_hidden_states = intermediary_hidden_states + self.mlp(
+ self.layernorm_2(intermediary_hidden_states)
+ )
+
+ return (intermediary_hidden_states,) + outputs
+
+
+@auto_docstring
+class BarkPreTrainedModel(PreTrainedModel):
+ config: BarkConfig
+ supports_gradient_checkpointing = False
+ _supports_flash_attn = True
+
+ @property
+ def device(self) -> torch.device:
+ """
+ `torch.device`: The device on which the module is (assuming that all the module parameters are on the same
+ device).
+ """
+
+ # if has _hf_hook, has been offloaded so the device has to be found in the hook
+ if not hasattr(self, "_hf_hook"):
+ return super().device
+ for module in self.modules():
+ if (
+ hasattr(module, "_hf_hook")
+ and hasattr(module._hf_hook, "execution_device")
+ and module._hf_hook.execution_device is not None
+ ):
+ return torch.device(module._hf_hook.execution_device)
+
+ return super().device
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, BarkSelfAttention):
+ if module.is_causal:
+ block_size = module.config.block_size
+ bias = torch.tril(torch.ones((block_size, block_size), dtype=bool)).view(1, 1, block_size, block_size)
+ init.copy_(module.bias, bias)
+
+
+# GPT2-like autoregressive model
+class BarkCausalModel(BarkPreTrainedModel, GenerationMixin):
+ config: BarkSubModelConfig
+ output_modalities = ("audio",)
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.config = config
+
+ # initialize as an autoregressive GPT-like model
+ self.input_embeds_layer = nn.Embedding(config.input_vocab_size, config.hidden_size)
+ self.position_embeds_layer = nn.Embedding(config.block_size, config.hidden_size)
+
+ self.drop = nn.Dropout(config.dropout)
+
+ self.layers = nn.ModuleList([BarkBlock(config, is_causal=True, layer_idx=i) for i in range(config.num_layers)])
+
+ self.layernorm_final = nn.LayerNorm(config.hidden_size, bias=config.bias)
+
+ self.lm_head = nn.Linear(config.hidden_size, config.output_vocab_size, bias=False)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ # NOTE: get_output_embeddings() must return None to prevent accidental weight tying.
+ # See e.g. https://github.com/huggingface/transformers/pull/39339#discussion_r2219126400
+ return None
+
+ def get_input_embeddings(self):
+ return self.input_embeds_layer
+
+ def set_input_embeddings(self, new_embeddings):
+ self.input_embeds_layer = new_embeddings
+
+ @deprecate_kwarg("input_embeds", version="5.6.0", new_name="inputs_embeds")
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ labels: torch.LongTensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor] | CausalLMOutputWithPast:
+ 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.return_dict
+
+ loss = None
+ if labels is not None:
+ raise NotImplementedError(
+ "Training is not implemented yet for Bark - ensure you do not pass `labels` to the model."
+ )
+
+ # Verify if inputs_embeds already exists
+ # then compute embeddings.
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+ elif inputs_embeds is not None and past_key_values is None:
+ # we want to return the inputs_embeds in priority so that it is in line with a weird hack
+ # of Bark which concatenate two bits of the inputs_embeds on the first forward pass of the semantic model
+ pass
+ elif input_ids is not None:
+ inputs_embeds = self.input_embeds_layer(input_ids) # token embeddings of shape (b, t, n_embd)
+ elif inputs_embeds is not None:
+ pass
+ else:
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+ input_shape = inputs_embeds.size()[:-1]
+ seq_length = input_shape[-1]
+
+ 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
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+ inputs_embeds = inputs_embeds.to(self.position_embeds_layer.weight.device)
+
+ if position_ids is None:
+ position_ids = torch.arange(
+ past_length,
+ seq_length + past_length,
+ dtype=torch.long,
+ device=self.position_embeds_layer.weight.device,
+ )
+ position_ids = position_ids.unsqueeze(0) # shape (1, seq_length)
+
+ position_ids = position_ids.to(self.position_embeds_layer.weight.device)
+ position_embeds = self.position_embeds_layer(position_ids) # position embeddings of shape (1, t, n_embd)
+
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ )
+
+ hidden_states = self.drop(inputs_embeds + position_embeds)
+ output_shape = input_shape + (hidden_states.size(-1),)
+
+ all_self_attentions = () if output_attentions else None
+ all_hidden_states = () if output_hidden_states else None
+
+ for i, block in enumerate(self.layers):
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ outputs = block(
+ hidden_states,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ )
+
+ hidden_states = outputs[0]
+
+ if output_attentions:
+ all_self_attentions = all_self_attentions + (outputs[1],)
+
+ hidden_states = self.layernorm_final(hidden_states)
+
+ hidden_states = hidden_states.view(output_shape)
+
+ # Add last hidden state
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ logits = self.lm_head(hidden_states)
+
+ if not return_dict:
+ return tuple(
+ v for v in [None, logits, past_key_values, all_hidden_states, all_self_attentions] if v is not None
+ )
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=past_key_values,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Bark semantic (or text) model. It shares the same architecture as the coarse model.
+ It is a GPT-2 like autoregressive model with a language modeling head on top.
+ """
+)
+class BarkSemanticModel(BarkCausalModel):
+ base_model_prefix = "semantic"
+ config: BarkSemanticConfig
+
+ def generate(
+ self,
+ input_ids: torch.Tensor,
+ semantic_generation_config: BarkSemanticGenerationConfig | None = None,
+ history_prompt: dict[str, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs,
+ ) -> torch.LongTensor:
+ """
+ Generates text semantic tokens from an input prompt and an additional optional `Bark` speaker prompt.
+
+ Args:
+ input_ids (`Optional[torch.Tensor]` of shape (batch_size, seq_len), *optional*):
+ Input ids, i.e tokenized input sentences. Will be truncated up to
+ semantic_generation_config.max_input_semantic_length tokens. Note that the output audios will be as
+ long as the longest generation among the batch.
+ semantic_generation_config (`BarkSemanticGenerationConfig`):
+ Generation config indicating how to generate the semantic tokens.
+ history_prompt (`Optional[dict[str,torch.Tensor]]`, *optional*):
+ Optional `Bark` speaker prompt.
+ attention_mask (`Optional[torch.Tensor]`, *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)
+ Returns:
+ torch.LongTensor: Output semantic tokens.
+ """
+ if semantic_generation_config is None:
+ raise ValueError("`semantic_generation_config` has to be provided")
+
+ batch_size = input_ids.shape[0]
+
+ max_input_semantic_length = semantic_generation_config.max_input_semantic_length
+
+ input_ids = input_ids + semantic_generation_config.text_encoding_offset
+
+ if attention_mask is not None:
+ input_ids = input_ids.masked_fill((1 - attention_mask).bool(), semantic_generation_config.text_pad_token)
+
+ if history_prompt is not None:
+ semantic_history = history_prompt["semantic_prompt"][-max_input_semantic_length:]
+ semantic_history = nn.functional.pad(
+ semantic_history,
+ (0, max_input_semantic_length - len(semantic_history)),
+ value=semantic_generation_config.semantic_pad_token,
+ mode="constant",
+ )
+ else:
+ semantic_history = torch.full(
+ (max_input_semantic_length,),
+ semantic_generation_config.semantic_pad_token,
+ device=self.device,
+ dtype=torch.int,
+ )
+
+ semantic_history = torch.repeat_interleave(semantic_history[None], batch_size, dim=0)
+
+ infer_array = torch.tensor(
+ [[semantic_generation_config.semantic_infer_token]] * batch_size, dtype=torch.int
+ ).to(self.device)
+
+ inputs_embeds = torch.cat(
+ [
+ self.input_embeds_layer(input_ids[:, :max_input_semantic_length])
+ + self.input_embeds_layer(semantic_history[:, : max_input_semantic_length + 1]),
+ self.input_embeds_layer(infer_array),
+ ],
+ dim=1,
+ )
+
+ tokens_to_suppress = list(
+ range(semantic_generation_config.semantic_vocab_size, semantic_generation_config.semantic_pad_token)
+ )
+ tokens_to_suppress.extend(
+ list(range(semantic_generation_config.semantic_pad_token + 1, self.config.output_vocab_size))
+ )
+
+ suppress_tokens_logits_processor = SuppressTokensLogitsProcessor(tokens_to_suppress, device=input_ids.device)
+
+ min_eos_p = kwargs.get("min_eos_p", semantic_generation_config.min_eos_p)
+ early_stopping_logits_processor = BarkEosPrioritizerLogitsProcessor(
+ eos_token_id=semantic_generation_config.eos_token_id, min_eos_p=min_eos_p, device=input_ids.device
+ )
+
+ # pass input_ids in order to stay consistent with the transformers generate method even though it is not used
+ # (except to get the input seq_len - that's why we keep the first 257 tokens)
+ semantic_output = super().generate(
+ torch.ones((batch_size, max_input_semantic_length + 1), dtype=torch.int, device=self.device),
+ inputs_embeds=inputs_embeds,
+ logits_processor=[suppress_tokens_logits_processor, early_stopping_logits_processor],
+ generation_config=semantic_generation_config,
+ **kwargs,
+ ) # size: 10048
+
+ # take the generated semantic tokens
+ if kwargs.get("return_dict_in_generate", False):
+ semantic_output = semantic_output.sequences[:, max_input_semantic_length + 1 :]
+ else:
+ semantic_output = semantic_output[:, max_input_semantic_length + 1 :]
+ return semantic_output
+
+
+@auto_docstring(
+ custom_intro="""
+ Bark coarse acoustics model.
+ It shares the same architecture as the semantic (or text) model. It is a GPT-2 like autoregressive model with a
+ language modeling head on top.
+ """
+)
+class BarkCoarseModel(BarkCausalModel):
+ base_model_prefix = "coarse_acoustics"
+ config: BarkCoarseConfig
+
+ def preprocess_histories(
+ self,
+ max_coarse_history: int,
+ semantic_to_coarse_ratio: int,
+ batch_size: int,
+ semantic_generation_config: int,
+ codebook_size: int,
+ history_prompt: dict[str, torch.Tensor] | None = None,
+ ):
+ """
+ Preprocess the optional `Bark` speaker prompts before `self.generate`.
+
+ Args:
+ max_coarse_history (`int`):
+ Maximum size of coarse tokens used.
+ semantic_to_coarse_ratio (`int`):
+ Ratio of semantic to coarse frequency
+ batch_size (`int`):
+ Batch size, i.e the number of samples.
+ semantic_generation_config (`BarkSemanticGenerationConfig`):
+ Generation config indicating how to generate the semantic tokens.
+ codebook_size (`int`):
+ Codebook channel size, i.e. the size of the output vocabulary per codebook channel.
+ history_prompt (`Optional[dict[str,torch.Tensor]]`):
+ Optional `Bark` speaker prompt.
+ Returns: Returns:
+ `tuple(torch.FloatTensor)`:
+ - **x_semantic_history** (`torch.FloatTensor` -- Processed semantic speaker prompt.
+ - **x_coarse_history** (`torch.FloatTensor`) -- Processed coarse speaker prompt.
+ """
+ if history_prompt is not None:
+ x_semantic_history = torch.repeat_interleave(history_prompt["semantic_prompt"][None], batch_size, dim=0)
+ # clone to avoid modifying history_prompt.coarse_prompt
+ x_coarse_history = history_prompt["coarse_prompt"].clone()
+
+ # offset x_coarse_history
+ if codebook_size is not None:
+ for n in range(1, x_coarse_history.shape[0]):
+ # offset
+ x_coarse_history[n, :] += codebook_size * n
+
+ # flatten x_coarse_history
+ x_coarse_history = torch.transpose(x_coarse_history, 0, 1).reshape(-1)
+
+ x_coarse_history = x_coarse_history + semantic_generation_config.semantic_vocab_size
+
+ x_coarse_history = torch.repeat_interleave(x_coarse_history[None], batch_size, dim=0)
+ # e.g: after SEMANTIC_VOCAB_SIZE (10000), 1024 tokens dedicated to first codebook, 1024 next tokens
+ # dedicated to second codebook.
+
+ max_semantic_history = int(np.floor(max_coarse_history / semantic_to_coarse_ratio))
+ # trim histories correctly
+ n_semantic_hist_provided = min(
+ [
+ max_semantic_history,
+ x_semantic_history.shape[1] - x_semantic_history.shape[1] % 2,
+ int(np.floor(x_coarse_history.shape[1] / semantic_to_coarse_ratio)),
+ ]
+ )
+
+ n_coarse_hist_provided = int(round(n_semantic_hist_provided * semantic_to_coarse_ratio))
+
+ x_semantic_history = x_semantic_history[:, -n_semantic_hist_provided:].int()
+ x_coarse_history = x_coarse_history[:, -n_coarse_hist_provided:].int()
+ # bit of a hack for time alignment (sounds better) - from Bark original implementation
+ x_coarse_history = x_coarse_history[:, :-2]
+
+ else:
+ # shape: (batch_size, 0)
+ x_semantic_history = torch.tensor([[]] * batch_size, dtype=torch.int, device=self.device)
+ x_coarse_history = torch.tensor([[]] * batch_size, dtype=torch.int, device=self.device)
+
+ return x_semantic_history, x_coarse_history
+
+ def generate(
+ self,
+ semantic_output: torch.Tensor,
+ semantic_generation_config: BarkSemanticGenerationConfig | None = None,
+ coarse_generation_config: BarkCoarseGenerationConfig | None = None,
+ codebook_size: int = 1024,
+ history_prompt: dict[str, torch.Tensor] | None = None,
+ return_output_lengths: bool | None = None,
+ **kwargs,
+ ) -> torch.LongTensor | tuple[torch.LongTensor, torch.LongTensor]:
+ """
+ Generates coarse acoustics tokens from input text semantic tokens and an additional optional `Bark` speaker
+ prompt.
+
+ Args:
+ semantic_output (`torch.Tensor` of shape (batch_size, seq_len), *optional*):
+ Input text semantic ids, i.e the output of `BarkSemanticModel.generate`.
+ semantic_generation_config (`BarkSemanticGenerationConfig`):
+ Generation config indicating how to generate the semantic tokens.
+ coarse_generation_config (`BarkCoarseGenerationConfig`):
+ Generation config indicating how to generate the coarse tokens.
+ codebook_size (`int`, *optional*, defaults to 1024):
+ Codebook channel size, i.e. the size of the output vocabulary per codebook channel.
+ history_prompt (`Optional[dict[str,torch.Tensor]]`, *optional*):
+ Optional `Bark` speaker prompt.
+ return_output_lengths (`bool`, *optional*):
+ Whether or not to return the output lengths. Useful when batching.
+ Returns:
+ By default:
+ torch.LongTensor: Output coarse acoustics tokens.
+ If `return_output_lengths=True`:
+ `Tuple(torch.Tensor, torch.Tensor): The output coarse acoustics tokens, and the length of each sample
+ of the batch.
+ """
+
+ if semantic_generation_config is None:
+ raise ValueError("`semantic_generation_config` has to be provided")
+
+ if coarse_generation_config is None:
+ raise ValueError("`coarse_generation_config` has to be provided")
+
+ max_coarse_input_length = coarse_generation_config.max_coarse_input_length
+ max_coarse_history = coarse_generation_config.max_coarse_history
+ sliding_window_len = coarse_generation_config.sliding_window_len
+
+ # replace semantic_pad_token (eos_tok and pad_tok here) with coarse_semantic_pad_token i.e the pad_token
+ # used in the next model
+ semantic_output.masked_fill_(
+ semantic_output == semantic_generation_config.semantic_pad_token,
+ coarse_generation_config.coarse_semantic_pad_token,
+ )
+
+ semantic_to_coarse_ratio = (
+ coarse_generation_config.coarse_rate_hz
+ / semantic_generation_config.semantic_rate_hz
+ * coarse_generation_config.n_coarse_codebooks
+ )
+ max_semantic_history = int(np.floor(max_coarse_history / semantic_to_coarse_ratio))
+
+ output_lengths = (semantic_output != coarse_generation_config.coarse_semantic_pad_token).sum(1)
+ output_lengths = torch.floor(
+ output_lengths * semantic_to_coarse_ratio / coarse_generation_config.n_coarse_codebooks
+ )
+ output_lengths = torch.round(output_lengths * coarse_generation_config.n_coarse_codebooks).int()
+
+ max_generated_len = torch.max(output_lengths).item()
+
+ batch_size = semantic_output.shape[0]
+
+ x_semantic_history, x_coarse = self.preprocess_histories(
+ history_prompt=history_prompt,
+ max_coarse_history=max_coarse_history,
+ semantic_to_coarse_ratio=semantic_to_coarse_ratio,
+ batch_size=batch_size,
+ semantic_generation_config=semantic_generation_config,
+ codebook_size=codebook_size,
+ )
+ base_semantic_idx = x_semantic_history.shape[1]
+
+ semantic_output = torch.hstack([x_semantic_history, semantic_output])
+
+ n_window_steps = int(np.ceil(max_generated_len / sliding_window_len))
+
+ total_generated_len = 0
+
+ len_coarse_history = x_coarse.shape[1]
+
+ for _ in range(n_window_steps):
+ semantic_idx = base_semantic_idx + int(round(total_generated_len / semantic_to_coarse_ratio))
+
+ # pad from right side
+ input_coarse = semantic_output[:, np.max([0, semantic_idx - max_semantic_history]) :]
+ input_coarse = input_coarse[:, :max_coarse_input_length]
+ input_coarse = F.pad(
+ input_coarse,
+ (0, max_coarse_input_length - input_coarse.shape[-1]),
+ "constant",
+ coarse_generation_config.coarse_semantic_pad_token,
+ )
+
+ input_coarse = torch.hstack(
+ [
+ input_coarse,
+ torch.tensor([[coarse_generation_config.coarse_infer_token]] * batch_size, device=self.device),
+ x_coarse[:, -max_coarse_history:],
+ ]
+ )
+
+ alternatingLogitsProcessor = AlternatingCodebooksLogitsProcessor(
+ input_coarse.shape[1],
+ semantic_generation_config.semantic_vocab_size,
+ codebook_size,
+ )
+
+ output_coarse = super().generate(
+ input_coarse,
+ logits_processor=[alternatingLogitsProcessor],
+ max_new_tokens=min(sliding_window_len, max_generated_len - total_generated_len),
+ generation_config=coarse_generation_config,
+ **kwargs,
+ )
+
+ input_coarse_len = input_coarse.shape[1]
+
+ if kwargs.get("return_dict_in_generate", False):
+ x_coarse = torch.hstack([x_coarse, output_coarse.sequences[:, input_coarse_len:]])
+ else:
+ x_coarse = torch.hstack([x_coarse, output_coarse[:, input_coarse_len:]])
+ total_generated_len = x_coarse.shape[1] - len_coarse_history
+
+ del output_coarse
+
+ coarse_output = x_coarse[:, len_coarse_history:]
+
+ if return_output_lengths:
+ return coarse_output, output_lengths
+
+ return coarse_output
+
+
+@auto_docstring(
+ custom_intro="""
+ Bark fine acoustics model. It is a non-causal GPT-like model with `config.n_codes_total` embedding layers and
+ language modeling heads, one for each codebook.
+ """
+)
+class BarkFineModel(BarkPreTrainedModel):
+ base_model_prefix = "fine_acoustics"
+ config: BarkFineConfig
+ main_input_name = "codebook_idx"
+
+ def __init__(self, config):
+ # non-causal gpt-like model with one embedding layer and one lm_head for each codebook of Encodec
+ super().__init__(config)
+ self.config = config
+ self._tied_weights_keys = {}
+ for i in range(self.config.n_codes_total - self.config.n_codes_given):
+ self._tied_weights_keys[f"lm_heads.{i}.weight"] = f"input_embeds_layers.{i + 1}.weight"
+
+ # initialize a modified non causal GPT-like model
+ # note that for there is one embedding layer and one lm_head for each codebook of Encodec
+ self.input_embeds_layers = nn.ModuleList(
+ [nn.Embedding(config.input_vocab_size, config.hidden_size) for _ in range(config.n_codes_total)]
+ )
+ self.position_embeds_layer = nn.Embedding(config.block_size, config.hidden_size)
+
+ self.drop = nn.Dropout(config.dropout)
+
+ self.layers = nn.ModuleList(
+ [BarkBlock(config, is_causal=False, layer_idx=i) for i in range(config.num_layers)]
+ )
+
+ self.layernorm_final = nn.LayerNorm(config.hidden_size)
+
+ self.lm_heads = nn.ModuleList(
+ [
+ nn.Linear(config.hidden_size, config.output_vocab_size, bias=False)
+ for _ in range(config.n_codes_given, config.n_codes_total)
+ ]
+ )
+ self.gradient_checkpointing = False
+ self.n_codes_total = config.n_codes_total
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ # one embedding layers for each codebook
+ return self.input_embeds_layers
+
+ def set_input_embeddings(self, new_embeddings):
+ # one embedding layers for each codebook
+ self.input_embeds_layers = new_embeddings
+
+ def get_output_embeddings(self):
+ # one lm_head for each codebook
+ return self.lm_heads
+
+ def set_output_embeddings(self, new_output_embeddings):
+ # one lm_head for each codebook
+ self.lm_heads = new_output_embeddings
+
+ def _resize_token_embeddings(self, new_num_tokens, pad_to_multiple_of=None, mean_resizing=True):
+ old_embeddings_list = self.get_input_embeddings()
+ new_embeddings_list = nn.ModuleList(
+ [
+ self._get_resized_embeddings(old_embeddings, new_num_tokens, pad_to_multiple_of, mean_resizing)
+ for old_embeddings in old_embeddings_list
+ ]
+ )
+ self.set_input_embeddings(new_embeddings_list)
+ new_num_tokens = new_embeddings_list[0].weight.shape[0]
+
+ # if word embeddings are not tied, make sure that lm head is resized as well
+ if self.get_output_embeddings() is not None and not self.config.tie_word_embeddings:
+ old_lm_head_list = self.get_output_embeddings()
+ new_lm_head_list = nn.ModuleList(
+ [self._get_resized_lm_head(old_lm_head, new_num_tokens) for old_lm_head in old_lm_head_list]
+ )
+ self.set_output_embeddings(new_lm_head_list)
+
+ return self.get_input_embeddings()
+
+ def resize_token_embeddings(
+ self,
+ new_num_tokens: int | None = None,
+ pad_to_multiple_of: int | None = None,
+ mean_resizing: bool = True,
+ ) -> nn.Embedding:
+ """
+ Resizes input token embeddings matrix of the model if `new_num_tokens != config.vocab_size`.
+
+ Takes care of tying weights embeddings afterwards if the model class has a `tie_weights()` method.
+
+ Arguments:
+ new_num_tokens (`int`, *optional*):
+ The number of new tokens in the embedding matrix. Increasing the size will add newly initialized
+ vectors at the end. Reducing the size will remove vectors from the end. If not provided or `None`, just
+ returns a pointer to the input tokens `torch.nn.Embedding` module of the model without doing anything.
+ pad_to_multiple_of (`int`, *optional*):
+ If set will pad the embedding matrix to a multiple of the provided value.
+
+ This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability
+ `>= 7.5` (Volta), or on TPUs which benefit from having sequence lengths be a multiple of 128. For more
+ details about this, or help on choosing the correct value for resizing, refer to this guide:
+ https://docs.nvidia.com/deeplearning/performance/dl-performance-matrix-multiplication/index.html#requirements-tc
+ mean_resizing (`bool`):
+ Whether to initialize the added embeddings from a multivariate normal distribution that has old embeddings' mean and
+ covariance or to initialize them with a normal distribution that has a mean of zero and std equals `config.initializer_range`.
+
+ Setting `mean_resizing` to `True` is useful when increasing the size of the embeddings of causal language models,
+ where the generated tokens' probabilities won't be affected by the added embeddings because initializing the new embeddings with the
+ old embeddings' mean will reduce the kl-divergence between the next token probability before and after adding the new embeddings.
+ Refer to this article for more information: https://nlp.stanford.edu/~johnhew/vocab-expansion.html
+
+ Return:
+ `torch.nn.Embedding`: Pointer to the input tokens Embeddings Module of the model.
+ """
+ model_embeds = self._resize_token_embeddings(new_num_tokens, pad_to_multiple_of, mean_resizing)
+ if new_num_tokens is None and pad_to_multiple_of is None:
+ return model_embeds
+
+ # Update base model and current model config
+ self.config.output_vocab_size = model_embeds[0].weight.shape[0]
+ self.config.vocab_size = model_embeds[0].weight.shape[0]
+ self.output_vocab_size = model_embeds[0].weight.shape[0]
+ self.vocab_size = model_embeds[0].weight.shape[0]
+
+ # Tie weights again if needed
+ self.tie_weights()
+
+ return model_embeds
+
+ @deprecate_kwarg("input_embeds", version="5.6.0", new_name="inputs_embeds")
+ @auto_docstring
+ def forward(
+ self,
+ codebook_idx: int, # an additional idx corresponding to the id of the codebook that will be predicted
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ labels: torch.LongTensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor] | MaskedLMOutput:
+ r"""
+ codebook_idx (`int`):
+ Index of the codebook that will be predicted.
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ NOT IMPLEMENTED YET.
+ """
+ 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.return_dict
+
+ loss = None
+ if labels is not None:
+ raise NotImplementedError("Training is not implemented yet")
+
+ if codebook_idx == 0:
+ raise ValueError("Cannot predict 0th codebook - 0th codebook should be predicted by the coarse model")
+
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+
+ if input_ids is None and inputs_embeds is None:
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+ if input_ids is not None:
+ # the input_embeddings are the sum of the j previous codebooks embeddings before
+ # the current codebook_idx codebook
+
+ # forward the GPT model itself
+ inputs_embeds = [
+ input_embeds_layer(input_ids[:, :, i]).unsqueeze(-1)
+ for i, input_embeds_layer in enumerate(self.input_embeds_layers)
+ ] # token embeddings of shape (b, t, n_embd)
+ inputs_embeds = torch.cat(inputs_embeds, dim=-1)
+ inputs_embeds = inputs_embeds[:, :, :, : codebook_idx + 1].sum(dim=-1)
+
+ input_shape = inputs_embeds.size()[:-1]
+ seq_length = input_shape[1]
+
+ inputs_embeds = inputs_embeds.to(self.position_embeds_layer.weight.device)
+
+ if position_ids is None:
+ position_ids = torch.arange(
+ 0, seq_length, dtype=torch.long, device=self.position_embeds_layer.weight.device
+ )
+ position_ids = position_ids.unsqueeze(0) # shape (1, seq_length)
+
+ position_ids = position_ids.to(self.position_embeds_layer.weight.device)
+ position_embeds = self.position_embeds_layer(position_ids) # position embeddings of shape (1, t, n_embd)
+
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ )
+
+ hidden_states = self.drop(inputs_embeds + position_embeds)
+ output_shape = input_shape + (hidden_states.size(-1),)
+
+ all_self_attentions = () if output_attentions else None
+ all_hidden_states = () if output_hidden_states else None
+
+ for i, block in enumerate(self.layers):
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ outputs = block(
+ hidden_states,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ )
+
+ hidden_states = outputs[0]
+
+ if output_attentions:
+ all_self_attentions = all_self_attentions + (outputs[1],)
+
+ hidden_states = self.layernorm_final(hidden_states)
+ hidden_states = hidden_states.view(output_shape)
+
+ # Add last hidden state
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ logits = self.lm_heads[codebook_idx - self.config.n_codes_given](hidden_states)
+
+ if not return_dict:
+ return tuple(v for v in [None, logits, all_hidden_states, all_self_attentions] if v is not None)
+
+ return MaskedLMOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ )
+
+ @torch.no_grad()
+ def generate(
+ self,
+ coarse_output: torch.Tensor,
+ semantic_generation_config: BarkSemanticGenerationConfig | None = None,
+ coarse_generation_config: BarkCoarseGenerationConfig | None = None,
+ fine_generation_config: BarkFineGenerationConfig = None,
+ codebook_size: int = 1024,
+ history_prompt: dict[str, torch.Tensor] | None = None,
+ **kwargs,
+ ) -> torch.LongTensor:
+ """
+ Generates fine acoustics tokens from input coarse acoustics tokens and an additional optional `Bark` speaker
+ prompt.
+
+ Args:
+ coarse_output (`torch.Tensor` of shape (batch_size, seq_len)):
+ Input coarse acoustics ids, i.e the output of `BarkCoarseModel.generate`.
+ semantic_generation_config (`BarkSemanticGenerationConfig`):
+ Generation config indicating how to generate the semantic tokens.
+ coarse_generation_config (`BarkCoarseGenerationConfig`):
+ Generation config indicating how to generate the coarse tokens.
+ fine_generation_config (`BarkFineGenerationConfig`):
+ Generation config indicating how to generate the fine tokens.
+ codebook_size (`int`, *optional*, defaults to 1024):
+ Codebook channel size, i.e. the size of the output vocabulary per codebook channel.
+ history_prompt (`Optional[dict[str,torch.Tensor]]`, *optional*):
+ Optional `Bark` speaker prompt.
+ Returns:
+ torch.LongTensor: Output fine acoustics tokens.
+ """
+ if semantic_generation_config is None:
+ raise ValueError("`semantic_generation_config` has to be provided")
+
+ if coarse_generation_config is None:
+ raise ValueError("`coarse_generation_config` has to be provided")
+
+ if fine_generation_config is None:
+ raise ValueError("`fine_generation_config` has to be provided")
+
+ # since we don't really use GenerationConfig through the fine model (autoencoder)
+ # and since only temperature is used from the classic GenerationConfig parameters
+ # manually impose the kwargs priority over the generation config
+ temperature = kwargs.get("temperature", fine_generation_config.temperature)
+
+ max_fine_history_length = fine_generation_config.max_fine_history_length
+ max_fine_input_length = fine_generation_config.max_fine_input_length
+
+ # shape: (batch, n_coarse_codebooks * seq_len)
+ # new_shape: (batch, seq_len, n_coarse_codebooks)
+ coarse_output = coarse_output.view(coarse_output.shape[0], -1, coarse_generation_config.n_coarse_codebooks)
+
+ # brings ids into the range [0, codebook_size -1]
+ coarse_output = torch.remainder(coarse_output - semantic_generation_config.semantic_vocab_size, codebook_size)
+ batch_size = coarse_output.shape[0]
+
+ if history_prompt is not None:
+ x_fine_history = torch.repeat_interleave(history_prompt["fine_prompt"].T[None], batch_size, dim=0)
+ # transpose to get to shape (seq_len, n_fine_codebooks)
+ else:
+ x_fine_history = None
+
+ n_coarse = coarse_generation_config.n_coarse_codebooks
+
+ # pad the last 6th codebooks
+ fine_input = F.pad(
+ coarse_output,
+ (0, fine_generation_config.n_fine_codebooks - n_coarse),
+ "constant",
+ codebook_size,
+ )
+
+ # prepend history if available (max max_fine_history_length)
+ if x_fine_history is not None:
+ fine_input = torch.cat([x_fine_history[:, -max_fine_history_length:, :], fine_input], dim=1)
+
+ # len of the fine_history that has been added to fine_input
+ n_history = x_fine_history[:, -max_fine_history_length:, :].shape[1]
+ else:
+ n_history = 0
+
+ n_remove_from_end = 0
+ # need to pad if too short (since non-causal model)
+ if fine_input.shape[1] < max_fine_input_length:
+ n_remove_from_end = max_fine_input_length - fine_input.shape[1]
+ fine_input = F.pad(fine_input, (0, 0, 0, n_remove_from_end), mode="constant", value=codebook_size)
+
+ # we can be lazy about fractional loop and just keep overwriting codebooks.
+ # seems that coarse_output.shape[1] - (max_fine_input_length - n_history) is equal to minus n_remove_from_end
+ # So if we needed to pad because too short, n_loops is always 1 (because n_remove_from_end > 0)
+ # If not, we loop over at least twice.
+
+ n_loops = (coarse_output.shape[1] - (max_fine_input_length - n_history)) / max_fine_history_length
+ n_loops = int(np.ceil(n_loops))
+ n_loops = max(0, n_loops) + 1
+
+ for n_outer in range(n_loops):
+ start_idx = min([n_outer * max_fine_history_length, fine_input.shape[1] - max_fine_input_length])
+
+ start_fill_idx = min(
+ [n_history + n_outer * max_fine_history_length, fine_input.shape[1] - max_fine_history_length]
+ )
+ rel_start_fill_idx = start_fill_idx - start_idx
+ input_buffer = fine_input[:, start_idx : start_idx + max_fine_input_length, :]
+ for n_inner in range(n_coarse, fine_generation_config.n_fine_codebooks):
+ logits = self.forward(n_inner, input_buffer).logits
+ if temperature is None or temperature == 1.0:
+ relevant_logits = logits[:, rel_start_fill_idx:, :codebook_size]
+ codebook_preds = torch.argmax(relevant_logits, -1)
+ else:
+ relevant_logits = logits[:, :, :codebook_size] / temperature
+ # apply softmax
+ probs = F.softmax(relevant_logits, dim=-1)[:, rel_start_fill_idx:max_fine_input_length]
+ # reshape to 2D: (batch_size, seq_len, codebook_size) -> (batch_size*seq_len, codebook_size)
+ probs = probs.reshape((-1, codebook_size))
+ # multinomial then reshape : (batch_size*seq_len)-> (batch_size,seq_len)
+ codebook_preds = torch.multinomial(probs, num_samples=1).view(batch_size, -1)
+ codebook_preds = codebook_preds.to(torch.int32)
+ input_buffer[:, rel_start_fill_idx:, n_inner] = codebook_preds
+ del logits, codebook_preds
+
+ # transfer into fine_input
+ for n_inner in range(n_coarse, fine_generation_config.n_fine_codebooks):
+ fine_input[
+ :, start_fill_idx : start_fill_idx + (max_fine_input_length - rel_start_fill_idx), n_inner
+ ] = input_buffer[:, rel_start_fill_idx:, n_inner]
+ del input_buffer
+
+ fine_input = fine_input.transpose(1, 2)[:, :, n_history:]
+ if n_remove_from_end > 0:
+ fine_input = fine_input[:, :, :-n_remove_from_end]
+
+ if fine_input.shape[-1] != coarse_output.shape[-2]:
+ raise ValueError("input and output should have the same seq_len")
+
+ return fine_input
+
+
+@auto_docstring(
+ custom_intro="""
+ The full Bark model, a text-to-speech model composed of 4 sub-models:
+ - [`BarkSemanticModel`] (also referred to as the 'text' model): a causal auto-regressive transformer model that
+ takes
+ as input tokenized text, and predicts semantic text tokens that capture the meaning of the text.
+ - [`BarkCoarseModel`] (also referred to as the 'coarse acoustics' model), also a causal autoregressive transformer,
+ that takes into input the results of the last model. It aims at regressing the first two audio codebooks necessary
+ to `encodec`.
+ - [`BarkFineModel`] (the 'fine acoustics' model), this time a non-causal autoencoder transformer, which iteratively
+ predicts the last codebooks based on the sum of the previous codebooks embeddings.
+ - having predicted all the codebook channels from the [`EncodecModel`], Bark uses it to decode the output audio
+ array.
+
+ It should be noted that each of the first three modules can support conditional speaker embeddings to condition the
+ output sound according to specific predefined voice.
+ """
+)
+class BarkModel(BarkPreTrainedModel, GenerationMixin):
+ config: BarkConfig
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.semantic = BarkSemanticModel(config.semantic_config)
+ self.coarse_acoustics = BarkCoarseModel(config.coarse_acoustics_config)
+ self.fine_acoustics = BarkFineModel(config.fine_acoustics_config)
+
+ self.codec_model = AutoModel.from_config(config.codec_config)
+
+ self.config = config
+
+ self.post_init()
+
+ @classmethod
+ def can_generate(cls) -> bool:
+ # Bark has a unique model structure, where the external class (`BarkModel`) doesn't need to inherit from
+ # `GenerationMixin` (it has a non-standard generation method), but one of the internal models do
+ # (`BarkSemanticModel`). This means that the base `can_generate()` will return `False`, but we need to
+ # override it so as to do `GenerationConfig` handling in multiple parts of the codebase.
+ return True
+
+ @property
+ def device(self) -> torch.device:
+ """
+ `torch.device`: The device on which the module is (assuming that all the module parameters are on the same
+ device).
+ """
+ # for bark_model, device must be verified on its sub-models
+ # if has _hf_hook, has been offloaded so the device has to be found in the hook
+ if not hasattr(self.semantic, "_hf_hook"):
+ return super().device
+ for module in self.semantic.modules():
+ if (
+ hasattr(module, "_hf_hook")
+ and hasattr(module._hf_hook, "execution_device")
+ and module._hf_hook.execution_device is not None
+ ):
+ return torch.device(module._hf_hook.execution_device)
+
+ def enable_cpu_offload(
+ self,
+ accelerator_id: int | None = 0,
+ **kwargs,
+ ):
+ r"""
+ Offloads all sub-models to CPU using accelerate, reducing memory usage with a low impact on performance. This
+ method moves one whole sub-model at a time to the accelerator when it is used, and the sub-model remains in accelerator until the next sub-model runs.
+
+ Args:
+ accelerator_id (`int`, *optional*, defaults to 0):
+ accelerator id on which the sub-models will be loaded and offloaded.
+ """
+ if is_accelerate_available():
+ from accelerate import cpu_offload_with_hook
+ else:
+ raise ImportError("`enable_model_cpu_offload` requires `accelerate`.")
+
+ device_type = "cuda"
+ if is_torch_accelerator_available():
+ device_type = torch.accelerator.current_accelerator().type
+ device = torch.device(f"{device_type}:{accelerator_id}")
+
+ torch_accelerator_module = getattr(torch, device_type)
+ if self.device.type != "cpu":
+ self.to("cpu")
+ torch_accelerator_module.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
+
+ # this layer is used outside the first forward pass of semantic so need to be loaded before semantic
+ self.semantic.input_embeds_layer, _ = cpu_offload_with_hook(self.semantic.input_embeds_layer, device)
+
+ hook = None
+ for cpu_offloaded_model in [
+ self.semantic,
+ self.coarse_acoustics,
+ self.fine_acoustics,
+ ]:
+ _, hook = cpu_offload_with_hook(cpu_offloaded_model, device, prev_module_hook=hook)
+
+ self.fine_acoustics_hook = hook
+
+ _, hook = cpu_offload_with_hook(self.codec_model, device, prev_module_hook=hook)
+
+ # We'll offload the last model manually.
+ self.codec_model_hook = hook
+
+ def codec_decode(self, fine_output, output_lengths=None):
+ """Turn quantized audio codes into audio array using encodec."""
+
+ fine_output = fine_output.transpose(0, 1)
+ emb = self.codec_model.quantizer.decode(fine_output)
+
+ if output_lengths is not None:
+ # encodec uses LSTMs which behaves differently with appended padding
+ # decoding with encodec takes around 0.1% of the total generation time
+ # to keep generation quality, we break batching
+ out = [sample[:, :l].unsqueeze(0) for (sample, l) in zip(emb, output_lengths)]
+ audio_arr = [self.codec_model.decoder(sample).squeeze() for sample in out]
+ else:
+ out = self.codec_model.decoder(emb)
+ audio_arr = out.squeeze(1) # squeeze the codebook dimension
+
+ return audio_arr
+
+ @torch.no_grad()
+ def generate(
+ self,
+ input_ids: torch.Tensor | None = None,
+ history_prompt: dict[str, torch.Tensor] | None = None,
+ return_output_lengths: bool | None = None,
+ **kwargs,
+ ) -> torch.LongTensor:
+ """
+ Generates audio from an input prompt and an additional optional `Bark` speaker prompt.
+
+ Args:
+ input_ids (`Optional[torch.Tensor]` of shape (batch_size, seq_len), *optional*):
+ Input ids. Will be truncated up to 256 tokens. Note that the output audios will be as long as the
+ longest generation among the batch.
+ history_prompt (`Optional[dict[str,torch.Tensor]]`, *optional*):
+ Optional `Bark` speaker prompt. Note that for now, this model takes only one speaker prompt per batch.
+ kwargs (*optional*): Remaining dictionary of keyword arguments. Keyword arguments are of two types:
+
+ - Without a prefix, they will be entered as `**kwargs` for the `generate` method of each sub-model.
+ - With a *semantic_*, *coarse_*, *fine_* prefix, they will be input for the `generate` method of the
+ semantic, coarse and fine respectively. It has the priority over the keywords without a prefix.
+
+ This means you can, for example, specify a generation strategy for all sub-models except one.
+ return_output_lengths (`bool`, *optional*):
+ Whether or not to return the waveform lengths. Useful when batching.
+ Returns:
+ By default:
+ - **audio_waveform** (`torch.Tensor` of shape (batch_size, seq_len)): Generated audio waveform.
+ When `return_output_lengths=True`:
+ Returns a tuple made of:
+ - **audio_waveform** (`torch.Tensor` of shape (batch_size, seq_len)): Generated audio waveform.
+ - **output_lengths** (`torch.Tensor` of shape (batch_size)): The length of each waveform in the batch
+ Example:
+
+ ```python
+ >>> from transformers import AutoProcessor, BarkModel
+
+ >>> processor = AutoProcessor.from_pretrained("suno/bark-small")
+ >>> model = BarkModel.from_pretrained("suno/bark-small")
+
+ >>> # To add a voice preset, you can pass `voice_preset` to `BarkProcessor.__call__(...)`
+ >>> voice_preset = "v2/en_speaker_6"
+
+ >>> inputs = processor("Hello, my dog is cute, I need him in my life", voice_preset=voice_preset)
+
+ >>> audio_array = model.generate(**inputs, semantic_max_new_tokens=100)
+ >>> audio_array = audio_array.cpu().numpy().squeeze()
+ ```
+ """
+ # TODO (joao):workaround until nested generation config is compatible with PreTrained Model
+ # todo: dict
+ semantic_generation_config = BarkSemanticGenerationConfig(**self.generation_config.semantic_config)
+ coarse_generation_config = BarkCoarseGenerationConfig(**self.generation_config.coarse_acoustics_config)
+ fine_generation_config = BarkFineGenerationConfig(**self.generation_config.fine_acoustics_config)
+
+ kwargs_semantic = {
+ # if "attention_mask" is set, it should not be passed to CoarseModel and FineModel
+ "attention_mask": kwargs.pop("attention_mask", None),
+ "min_eos_p": kwargs.pop("min_eos_p", None),
+ }
+ kwargs_coarse = {}
+ kwargs_fine = {}
+ for key, value in kwargs.items():
+ if key.startswith("semantic_"):
+ key = key[len("semantic_") :]
+ kwargs_semantic[key] = value
+ elif key.startswith("coarse_"):
+ key = key[len("coarse_") :]
+ kwargs_coarse[key] = value
+ elif key.startswith("fine_"):
+ key = key[len("fine_") :]
+ kwargs_fine[key] = value
+ else:
+ # If the key is already in a specific config, then it's been set with a
+ # submodules specific value and we don't override
+ if key not in kwargs_semantic:
+ kwargs_semantic[key] = value
+ if key not in kwargs_coarse:
+ kwargs_coarse[key] = value
+ if key not in kwargs_fine:
+ kwargs_fine[key] = value
+
+ # 1. Generate from the semantic model
+ if "generation_config" in kwargs_semantic:
+ kwargs_semantic.pop("generation_config")
+ semantic_output = self.semantic.generate(
+ input_ids,
+ history_prompt=history_prompt,
+ semantic_generation_config=semantic_generation_config,
+ **kwargs_semantic,
+ )
+
+ # 2. Generate from the coarse model
+ if "generation_config" in kwargs_coarse:
+ kwargs_coarse.pop("generation_config")
+ coarse_output = self.coarse_acoustics.generate(
+ semantic_output,
+ history_prompt=history_prompt,
+ semantic_generation_config=semantic_generation_config,
+ coarse_generation_config=coarse_generation_config,
+ codebook_size=self.generation_config.codebook_size,
+ return_output_lengths=return_output_lengths,
+ **kwargs_coarse,
+ )
+
+ output_lengths = None
+ if return_output_lengths:
+ coarse_output, output_lengths = coarse_output
+ # (batch_size, seq_len*coarse_codebooks) -> (batch_size, seq_len)
+ output_lengths = output_lengths // coarse_generation_config.n_coarse_codebooks
+
+ # 3. "generate" from the fine model
+ if "generation_config" in kwargs_fine:
+ kwargs_fine.pop("generation_config")
+ output = self.fine_acoustics.generate(
+ coarse_output,
+ history_prompt=history_prompt,
+ semantic_generation_config=semantic_generation_config,
+ coarse_generation_config=coarse_generation_config,
+ fine_generation_config=fine_generation_config,
+ codebook_size=self.generation_config.codebook_size,
+ **kwargs_fine,
+ )
+
+ if getattr(self, "fine_acoustics_hook", None) is not None:
+ # Manually offload fine_acoustics to CPU
+ # and load codec_model to GPU
+ # since bark doesn't use codec_model forward pass
+ self.fine_acoustics_hook.offload()
+ self.codec_model = self.codec_model.to(self.device)
+
+ # 4. Decode the output and generate audio array
+ audio = self.codec_decode(output, output_lengths)
+
+ if getattr(self, "codec_model_hook", None) is not None:
+ # Offload codec_model to CPU
+ self.codec_model_hook.offload()
+
+ if return_output_lengths:
+ output_lengths = [len(sample) for sample in audio]
+ audio = nn.utils.rnn.pad_sequence(audio, batch_first=True, padding_value=0)
+ return audio, output_lengths
+
+ return audio
+
+
+__all__ = [
+ "BarkFineModel",
+ "BarkSemanticModel",
+ "BarkCoarseModel",
+ "BarkModel",
+ "BarkPreTrainedModel",
+ "BarkCausalModel",
+]
diff --git a/third_party/transformers/src/transformers/models/bark/processing_bark.py b/third_party/transformers/src/transformers/models/bark/processing_bark.py
new file mode 100644
index 0000000000000000000000000000000000000000..e9b360c49279a4b4c07f31508973e5eea26dc6d8
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bark/processing_bark.py
@@ -0,0 +1,314 @@
+# Copyright 2023 The Suno AI Authors and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""
+Processor class for Bark
+"""
+
+import json
+import os
+
+import numpy as np
+
+from ...feature_extraction_utils import BatchFeature
+from ...processing_utils import ProcessorMixin
+from ...tokenization_utils_base import BatchEncoding
+from ...utils import auto_docstring, logging
+from ...utils.hub import cached_file
+from ..auto import AutoTokenizer
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring
+class BarkProcessor(ProcessorMixin):
+ preset_shape = {
+ "semantic_prompt": 1, # 1D array of shape (X,)
+ "coarse_prompt": 2, # 2D array of shape (2,X)
+ "fine_prompt": 2, # 2D array of shape (8,X)
+ }
+
+ def __init__(self, tokenizer, speaker_embeddings=None):
+ r"""
+ speaker_embeddings (`dict[dict[str]]`, *optional*):
+ Optional nested speaker embeddings dictionary. The first level contains voice preset names (e.g
+ `"en_speaker_4"`). The second level contains `"semantic_prompt"`, `"coarse_prompt"` and `"fine_prompt"`
+ embeddings. The values correspond to the path of the corresponding `np.ndarray`. See
+ [here](https://suno-ai.notion.site/8b8e8749ed514b0cbf3f699013548683?v=bc67cff786b04b50b3ceb756fd05f68c) for
+ a list of `voice_preset_names`.
+ """
+ super().__init__(tokenizer)
+
+ self.speaker_embeddings = speaker_embeddings
+
+ @classmethod
+ def from_pretrained(
+ cls, pretrained_processor_name_or_path, speaker_embeddings_dict_path="speaker_embeddings_path.json", **kwargs
+ ):
+ r"""
+ Instantiate a Bark processor associated with a pretrained model.
+
+ Args:
+ pretrained_model_name_or_path (`str` or `os.PathLike`):
+ This can be either:
+
+ - a string, the *model id* of a pretrained [`BarkProcessor`] hosted inside a model repo on
+ huggingface.co.
+ - a path to a *directory* containing a processor saved using the [`~BarkProcessor.save_pretrained`]
+ method, e.g., `./my_model_directory/`.
+ speaker_embeddings_dict_path (`str`, *optional*, defaults to `"speaker_embeddings_path.json"`):
+ The name of the `.json` file containing the speaker_embeddings dictionary located in
+ `pretrained_model_name_or_path`. If `None`, no speaker_embeddings is loaded.
+ **kwargs
+ Additional keyword arguments passed along to both
+ [`~tokenization_utils_base.PreTrainedTokenizer.from_pretrained`].
+ """
+ token = kwargs.get("token")
+ if speaker_embeddings_dict_path is not None:
+ speaker_embeddings_path = cached_file(
+ pretrained_processor_name_or_path,
+ speaker_embeddings_dict_path,
+ subfolder=kwargs.pop("subfolder", None),
+ cache_dir=kwargs.pop("cache_dir", None),
+ force_download=kwargs.pop("force_download", False),
+ proxies=kwargs.pop("proxies", None),
+ local_files_only=kwargs.pop("local_files_only", False),
+ token=token,
+ revision=kwargs.pop("revision", None),
+ _raise_exceptions_for_gated_repo=False,
+ _raise_exceptions_for_missing_entries=False,
+ _raise_exceptions_for_connection_errors=False,
+ )
+ if speaker_embeddings_path is None:
+ logger.warning(
+ f"""`{os.path.join(pretrained_processor_name_or_path, speaker_embeddings_dict_path)}` does not exists
+ , no preloaded speaker embeddings will be used - Make sure to provide a correct path to the json
+ dictionary if wanted, otherwise set `speaker_embeddings_dict_path=None`."""
+ )
+ speaker_embeddings = None
+ else:
+ with open(speaker_embeddings_path) as speaker_embeddings_json:
+ speaker_embeddings = json.load(speaker_embeddings_json)
+ else:
+ speaker_embeddings = None
+
+ if speaker_embeddings is not None:
+ if "repo_or_path" in speaker_embeddings:
+ speaker_embeddings["repo_or_path"] = pretrained_processor_name_or_path
+ tokenizer = AutoTokenizer.from_pretrained(pretrained_processor_name_or_path, **kwargs)
+
+ return cls(tokenizer=tokenizer, speaker_embeddings=speaker_embeddings)
+
+ def save_pretrained(
+ self,
+ save_directory,
+ speaker_embeddings_dict_path="speaker_embeddings_path.json",
+ speaker_embeddings_directory="speaker_embeddings",
+ push_to_hub: bool = False,
+ **kwargs,
+ ):
+ """
+ Saves the attributes of this processor (tokenizer...) in the specified directory so that it can be reloaded
+ using the [`~BarkProcessor.from_pretrained`] method.
+
+ Args:
+ save_directory (`str` or `os.PathLike`):
+ Directory where the tokenizer files and the speaker embeddings will be saved (directory will be created
+ if it does not exist).
+ speaker_embeddings_dict_path (`str`, *optional*, defaults to `"speaker_embeddings_path.json"`):
+ The name of the `.json` file that will contains the speaker_embeddings nested path dictionary, if it
+ exists, and that will be located in `pretrained_model_name_or_path/speaker_embeddings_directory`.
+ speaker_embeddings_directory (`str`, *optional*, defaults to `"speaker_embeddings/"`):
+ The name of the folder in which the speaker_embeddings arrays will be saved.
+ push_to_hub (`bool`, *optional*, defaults to `False`):
+ Whether or not to push your model to the Hugging Face model hub after saving it. You can specify the
+ repository you want to push to with `repo_id` (will default to the name of `save_directory` in your
+ namespace).
+ kwargs:
+ Additional key word arguments passed along to the [`~utils.PushToHubMixin.push_to_hub`] method.
+ """
+ if self.speaker_embeddings is not None:
+ os.makedirs(os.path.join(save_directory, speaker_embeddings_directory, "v2"), exist_ok=True)
+
+ embeddings_dict = {}
+
+ embeddings_dict["repo_or_path"] = save_directory
+
+ for prompt_key in self.available_voice_presets:
+ voice_preset = self._load_voice_preset(prompt_key)
+
+ tmp_dict = {}
+ for key in self.speaker_embeddings[prompt_key]:
+ np.save(
+ os.path.join(
+ embeddings_dict["repo_or_path"], speaker_embeddings_directory, f"{prompt_key}_{key}"
+ ),
+ voice_preset[key],
+ allow_pickle=False,
+ )
+ tmp_dict[key] = os.path.join(speaker_embeddings_directory, f"{prompt_key}_{key}.npy")
+
+ embeddings_dict[prompt_key] = tmp_dict
+
+ with open(os.path.join(save_directory, speaker_embeddings_dict_path), "w") as fp:
+ json.dump(embeddings_dict, fp)
+
+ super().save_pretrained(save_directory, push_to_hub, **kwargs)
+
+ def _load_voice_preset(self, voice_preset: str | None = None, **kwargs):
+ voice_preset_paths = self.speaker_embeddings[voice_preset]
+
+ voice_preset_dict = {}
+ token = kwargs.get("token")
+ for key in ["semantic_prompt", "coarse_prompt", "fine_prompt"]:
+ if key not in voice_preset_paths:
+ raise ValueError(
+ f"Voice preset unrecognized, missing {key} as a key in self.speaker_embeddings[{voice_preset}]."
+ )
+
+ path = cached_file(
+ self.speaker_embeddings.get("repo_or_path", "/"),
+ voice_preset_paths[key],
+ subfolder=kwargs.pop("subfolder", None),
+ cache_dir=kwargs.pop("cache_dir", None),
+ force_download=kwargs.pop("force_download", False),
+ proxies=kwargs.pop("proxies", None),
+ local_files_only=kwargs.pop("local_files_only", False),
+ token=token,
+ revision=kwargs.pop("revision", None),
+ _raise_exceptions_for_gated_repo=False,
+ _raise_exceptions_for_missing_entries=False,
+ _raise_exceptions_for_connection_errors=False,
+ )
+ if path is None:
+ raise ValueError(
+ f"""`{os.path.join(self.speaker_embeddings.get("repo_or_path", "/"), voice_preset_paths[key])}` does not exists
+ , no preloaded voice preset will be used - Make sure to provide correct paths to the {voice_preset}
+ embeddings."""
+ )
+
+ voice_preset_dict[key] = np.load(path)
+
+ return voice_preset_dict
+
+ def _validate_voice_preset_dict(self, voice_preset: dict | None = None):
+ for key in ["semantic_prompt", "coarse_prompt", "fine_prompt"]:
+ if key not in voice_preset:
+ raise ValueError(f"Voice preset unrecognized, missing {key} as a key.")
+
+ if not isinstance(voice_preset[key], np.ndarray):
+ raise TypeError(f"{key} voice preset must be a {str(self.preset_shape[key])}D ndarray.")
+
+ if len(voice_preset[key].shape) != self.preset_shape[key]:
+ raise ValueError(f"{key} voice preset must be a {str(self.preset_shape[key])}D ndarray.")
+
+ @property
+ def available_voice_presets(self) -> list:
+ """
+ Returns a list of available voice presets.
+
+ Returns:
+ `list[str]`: A list of voice preset names.
+ """
+ if self.speaker_embeddings is None:
+ return []
+
+ voice_presets = list(self.speaker_embeddings.keys())
+ if "repo_or_path" in voice_presets:
+ voice_presets.remove("repo_or_path")
+ return voice_presets
+
+ def _verify_speaker_embeddings(self, remove_unavailable: bool = True):
+ # check which actually downloaded properly / are available
+ unavailable_keys = []
+ if self.speaker_embeddings is not None:
+ for voice_preset in self.available_voice_presets:
+ try:
+ voice_preset_dict = self._load_voice_preset(voice_preset)
+ except ValueError:
+ # error from `_load_voice_preset` of path not existing
+ unavailable_keys.append(voice_preset)
+ continue
+ self._validate_voice_preset_dict(voice_preset_dict)
+
+ if unavailable_keys:
+ logger.warning(
+ f"The following {len(unavailable_keys)} speaker embeddings are not available: {unavailable_keys} "
+ "If you would like to use them, please check the paths or try downloading them again."
+ )
+
+ if remove_unavailable:
+ for voice_preset in unavailable_keys:
+ del self.speaker_embeddings[voice_preset]
+
+ @auto_docstring
+ def __call__(
+ self,
+ text=None,
+ voice_preset=None,
+ return_tensors="pt",
+ max_length=256,
+ add_special_tokens=False,
+ return_attention_mask=True,
+ return_token_type_ids=False,
+ **kwargs,
+ ) -> BatchEncoding:
+ r"""
+ voice_preset (`str`, `dict[np.ndarray]`):
+ The voice preset, i.e the speaker embeddings. It can either be a valid voice_preset name, e.g
+ `"en_speaker_1"`, or directly a dictionary of `np.ndarray` embeddings for each submodel of `Bark`. Or
+ it can be a valid file name of a local `.npz` single voice preset containing the keys
+ `"semantic_prompt"`, `"coarse_prompt"` and `"fine_prompt"`.
+
+ Returns:
+ [`BatchEncoding`]: A [`BatchEncoding`] object containing the output of the `tokenizer`.
+ If a voice preset is provided, the returned object will include a `"history_prompt"` key
+ containing a [`BatchFeature`], i.e the voice preset with the right tensors type.
+ """
+ if voice_preset is not None and not isinstance(voice_preset, dict):
+ if (
+ isinstance(voice_preset, str)
+ and self.speaker_embeddings is not None
+ and voice_preset in self.speaker_embeddings
+ ):
+ voice_preset = self._load_voice_preset(voice_preset)
+
+ else:
+ if isinstance(voice_preset, str) and not voice_preset.endswith(".npz"):
+ voice_preset = voice_preset + ".npz"
+
+ voice_preset = np.load(voice_preset)
+
+ if voice_preset is not None:
+ self._validate_voice_preset_dict(voice_preset, **kwargs)
+ voice_preset = BatchFeature(data=voice_preset, tensor_type=return_tensors)
+
+ encoded_text = self.tokenizer(
+ text,
+ return_tensors=return_tensors,
+ padding="max_length",
+ max_length=max_length,
+ return_attention_mask=return_attention_mask,
+ return_token_type_ids=return_token_type_ids,
+ add_special_tokens=add_special_tokens,
+ **kwargs,
+ )
+
+ if voice_preset is not None:
+ encoded_text["history_prompt"] = voice_preset
+
+ return encoded_text
+
+
+__all__ = ["BarkProcessor"]
diff --git a/third_party/transformers/src/transformers/models/bart/__init__.py b/third_party/transformers/src/transformers/models/bart/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..3cf9c34518954100c569a8c3a78e652a48b2a9f6
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bart/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from ..roberta.tokenization_roberta import RobertaTokenizer as BartTokenizer
+ from .configuration_bart import *
+ from .modeling_bart import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/bart/configuration_bart.py b/third_party/transformers/src/transformers/models/bart/configuration_bart.py
new file mode 100644
index 0000000000000000000000000000000000000000..ed6ea32d64f421b132850c9b8f219741bce29209
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bart/configuration_bart.py
@@ -0,0 +1,86 @@
+# Copyright 2021 The Fairseq Authors and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""BART model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="facebook/bart-large")
+@strict
+class BartConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import BartConfig, BartModel
+
+ >>> # Initializing a BART facebook/bart-large style configuration
+ >>> configuration = BartConfig()
+
+ >>> # Initializing a model (with random weights) from the facebook/bart-large style configuration
+ >>> model = BartModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "bart"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ attribute_map = {
+ "num_attention_heads": "encoder_attention_heads",
+ "hidden_size": "d_model",
+ "num_hidden_layers": "encoder_layers",
+ }
+
+ vocab_size: int = 50265
+ max_position_embeddings: int = 1024
+ encoder_layers: int | None = 12
+ encoder_ffn_dim: int | None = 4096
+ encoder_attention_heads: int | None = 16
+ decoder_layers: int | None = 12
+ decoder_ffn_dim: int | None = 4096
+ decoder_attention_heads: int | None = 16
+ encoder_layerdrop: float | None = 0.0
+ decoder_layerdrop: float | None = 0.0
+ activation_function: str | None = "gelu"
+ d_model: int | None = 1024
+ dropout: float | int | None = 0.1
+ attention_dropout: float | int | None = 0.0
+ activation_dropout: float | int | None = 0.0
+ init_std: float | None = 0.02
+ classifier_dropout: float | int | None = 0.0
+ scale_embedding: bool | None = False
+ use_cache: bool = True
+ pad_token_id: int | None = 1
+ bos_token_id: int | None = 0
+ eos_token_id: int | list[int] | None = 2
+ is_encoder_decoder: bool | None = True
+ decoder_start_token_id: int | None = 2
+ forced_eos_token_id: int | list[int] | None = 2
+ is_decoder: bool | None = False
+ tie_word_embeddings: bool = True
+
+ def __post_init__(self, **kwargs):
+ # Set the default `num_labels` only if `id2label` is not
+ # yet set, i.e. user didn't pass `id2label/lable2id` in kwargs
+ if self.id2label is None:
+ self.num_labels = kwargs.pop("num_labels", 3)
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["BartConfig"]
diff --git a/third_party/transformers/src/transformers/models/bart/convert_bart_original_pytorch_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/bart/convert_bart_original_pytorch_checkpoint_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..ebb295e5b046575c3032ef75b87194613a42a648
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bart/convert_bart_original_pytorch_checkpoint_to_pytorch.py
@@ -0,0 +1,155 @@
+# Copyright 2020 The HuggingFace Inc. team.
+#
+# 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.
+"""Convert BART checkpoint."""
+
+import argparse
+import os
+from pathlib import Path
+
+import fairseq
+import torch
+from packaging import version
+from torch import nn
+
+from transformers import (
+ BartConfig,
+ BartForConditionalGeneration,
+ BartForSequenceClassification,
+ BartModel,
+ BartTokenizer,
+)
+from transformers.utils import logging
+
+
+FAIRSEQ_MODELS = ["bart.large", "bart.large.mnli", "bart.large.cnn", "bart_xsum/model.pt"]
+extra_arch = {"bart.large": BartModel, "bart.large.mnli": BartForSequenceClassification}
+if version.parse(fairseq.__version__) < version.parse("0.9.0"):
+ raise Exception("requires fairseq >= 0.9.0")
+
+
+logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+SAMPLE_TEXT = " Hello world! cécé herlolip"
+
+mnli_rename_keys = [
+ ("model.classification_heads.mnli.dense.weight", "classification_head.dense.weight"),
+ ("model.classification_heads.mnli.dense.bias", "classification_head.dense.bias"),
+ ("model.classification_heads.mnli.out_proj.weight", "classification_head.out_proj.weight"),
+ ("model.classification_heads.mnli.out_proj.bias", "classification_head.out_proj.bias"),
+]
+
+
+def remove_ignore_keys_(state_dict):
+ ignore_keys = [
+ "encoder.version",
+ "decoder.version",
+ "model.encoder.version",
+ "model.decoder.version",
+ "_float_tensor",
+ ]
+ for k in ignore_keys:
+ state_dict.pop(k, None)
+
+
+def rename_key(dct, old, new):
+ val = dct.pop(old)
+ dct[new] = val
+
+
+def load_xsum_checkpoint(checkpoint_path):
+ """Checkpoint path should end in model.pt"""
+ sd = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+ hub_interface = torch.hub.load("pytorch/fairseq", "bart.large.cnn").eval()
+ hub_interface.model.load_state_dict(sd["model"])
+ return hub_interface
+
+
+def make_linear_from_emb(emb):
+ vocab_size, emb_size = emb.weight.shape
+ lin_layer = nn.Linear(vocab_size, emb_size, bias=False)
+ lin_layer.weight.data = emb.weight.data
+ return lin_layer
+
+
+@torch.no_grad()
+def convert_bart_checkpoint(checkpoint_path, pytorch_dump_folder_path, hf_checkpoint_name=None):
+ """
+ Copy/paste/tweak model's weights to our BERT structure.
+ """
+ if not os.path.exists(checkpoint_path):
+ bart = torch.hub.load("pytorch/fairseq", checkpoint_path).eval()
+ else:
+ bart = load_xsum_checkpoint(checkpoint_path)
+
+ bart.model.upgrade_state_dict(bart.model.state_dict())
+ if hf_checkpoint_name is None:
+ hf_checkpoint_name = checkpoint_path.replace(".", "-")
+ config = BartConfig.from_pretrained(hf_checkpoint_name)
+ tokens = bart.encode(SAMPLE_TEXT).unsqueeze(0)
+ tokens2 = BartTokenizer.from_pretrained(hf_checkpoint_name).encode(SAMPLE_TEXT, return_tensors="pt").unsqueeze(0)
+ if not torch.eq(tokens, tokens2).all():
+ raise ValueError(
+ f"converted tokenizer and pretrained tokenizer returned different output: {tokens} != {tokens2}"
+ )
+
+ if checkpoint_path == "bart.large.mnli":
+ state_dict = bart.state_dict()
+ remove_ignore_keys_(state_dict)
+ state_dict["model.shared.weight"] = state_dict["model.decoder.embed_tokens.weight"]
+ for src, dest in mnli_rename_keys:
+ rename_key(state_dict, src, dest)
+ model = BartForSequenceClassification(config).eval()
+ model.load_state_dict(state_dict)
+ fairseq_output = bart.predict("mnli", tokens, return_logits=True)
+ new_model_outputs = model(tokens)[0] # logits
+ else: # no classification heads to worry about
+ state_dict = bart.model.state_dict()
+ remove_ignore_keys_(state_dict)
+ state_dict["shared.weight"] = state_dict["decoder.embed_tokens.weight"]
+ fairseq_output = bart.extract_features(tokens)
+ if hf_checkpoint_name == "facebook/bart-large":
+ model = BartModel(config).eval()
+ model.load_state_dict(state_dict)
+ new_model_outputs = model(tokens).model[0]
+ else:
+ model = BartForConditionalGeneration(config).eval() # an existing summarization ckpt
+ model.model.load_state_dict(state_dict)
+ if hasattr(model, "lm_head"):
+ model.lm_head = make_linear_from_emb(model.model.shared)
+ new_model_outputs = model.model(tokens)[0]
+
+ # Check results
+ if fairseq_output.shape != new_model_outputs.shape:
+ raise ValueError(
+ f"`fairseq_output` shape and `new_model_output` shape are different: {fairseq_output.shape=}, {new_model_outputs.shape}"
+ )
+ if (fairseq_output != new_model_outputs).any().item():
+ raise ValueError("Some values in `fairseq_output` are different from `new_model_outputs`")
+ Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
+ model.save_pretrained(pytorch_dump_folder_path)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ # Required parameters
+ parser.add_argument(
+ "fairseq_path", type=str, help="bart.large, bart.large.cnn or a path to a model.pt on local filesystem."
+ )
+ parser.add_argument("pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
+ parser.add_argument(
+ "--hf_config", default=None, type=str, help="Which huggingface architecture to use: bart-large-xsum"
+ )
+ args = parser.parse_args()
+ convert_bart_checkpoint(args.fairseq_path, args.pytorch_dump_folder_path, hf_checkpoint_name=args.hf_config)
diff --git a/third_party/transformers/src/transformers/models/bart/modeling_bart.py b/third_party/transformers/src/transformers/models/bart/modeling_bart.py
new file mode 100644
index 0000000000000000000000000000000000000000..d5f3e73b1e4bae6d55db14020650549704ca78ac
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bart/modeling_bart.py
@@ -0,0 +1,1322 @@
+# Copyright 2021 The Fairseq Authors and The HuggingFace Inc. team. All rights reserved.
+#
+# 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 BART model."""
+
+import math
+import warnings
+from collections.abc import Callable
+
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_bidirectional_mask, create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPastAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions,
+ Seq2SeqLMOutput,
+ Seq2SeqModelOutput,
+ Seq2SeqQuestionAnsweringModelOutput,
+ Seq2SeqSequenceClassifierOutput,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import (
+ TransformersKwargs,
+ auto_docstring,
+ can_return_tuple,
+ is_torchdynamo_compiling,
+ logging,
+ torch_compilable_check,
+)
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from .configuration_bart import BartConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
+ """
+ Shift input ids one token to the right.
+ """
+ shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+ shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
+ shifted_input_ids[:, 0] = decoder_start_token_id
+
+ if pad_token_id is None:
+ raise ValueError("self.model.config.pad_token_id has to be defined.")
+ # replace possible -100 values in labels by `pad_token_id`
+ shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+ return shifted_input_ids
+
+
+class BartLearnedPositionalEmbedding(nn.Embedding):
+ """
+ This module learns positional embeddings up to a fixed maximum size.
+ """
+
+ def __init__(self, num_embeddings: int, embedding_dim: int):
+ # Bart is set up so that if padding_idx is specified then offset the embedding ids by 2
+ # and adjust num_embeddings appropriately. Other models don't have this hack
+ self.offset = 2
+ super().__init__(num_embeddings + self.offset, embedding_dim)
+
+ def forward(
+ self, input_ids: torch.Tensor, past_key_values_length: int = 0, position_ids: torch.Tensor | None = None
+ ):
+ """`input_ids' shape is expected to be [bsz x seqlen]."""
+
+ if position_ids is None:
+ bsz, seq_len = input_ids.shape[:2]
+ position_ids = torch.arange(
+ past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device
+ ).expand(bsz, -1)
+ else:
+ position_ids = position_ids.unsqueeze(0)
+
+ return super().forward(position_ids + self.offset)
+
+
+class BartScaledWordEmbedding(nn.Embedding):
+ """
+ This module overrides nn.Embeddings' forward by multiplying with embeddings scale.
+ """
+
+ def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: float | None = 1.0):
+ super().__init__(num_embeddings, embedding_dim, padding_idx)
+ self.embed_scale = embed_scale
+
+ def forward(self, input_ids: torch.Tensor):
+ return super().forward(input_ids) * self.embed_scale
+
+
+# Copied from transformers.models.bert.modeling_bert.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float | None = None,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ if scaling is None:
+ scaling = query.size(-1) ** -0.5
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class BartAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ dropout: float = 0.0,
+ is_decoder: bool = False,
+ bias: bool = True,
+ is_causal: bool = False,
+ config: BartConfig | None = None,
+ layer_idx: int | None = None,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.dropout = dropout
+ self.head_dim = embed_dim // num_heads
+ self.config = config
+
+ if (self.head_dim * num_heads) != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+ f" and `num_heads`: {num_heads})."
+ )
+ self.scaling = self.head_dim**-0.5
+ self.is_decoder = is_decoder
+ self.is_causal = is_causal
+ self.layer_idx = layer_idx
+ if layer_idx is None and self.is_decoder:
+ logger.warning_once(
+ f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and "
+ "will lead to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class."
+ )
+
+ self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ key_value_states: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ # TODO: we need a refactor so that the different attention modules can get their specific kwargs
+ # ATM, we have mixed things encoder, decoder, and encoder-decoder attn
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Input shape: Batch x Time x Channel"""
+
+ # if key_value_states are provided this layer is used as a cross-attention layer
+ # for the decoder
+ is_cross_attention = key_value_states is not None
+
+ # determine input shapes
+ bsz, tgt_len = hidden_states.shape[:-1]
+ src_len = key_value_states.shape[1] if is_cross_attention else tgt_len
+
+ q_input_shape = (bsz, tgt_len, -1, self.head_dim)
+ kv_input_shape = (bsz, src_len, -1, self.head_dim)
+
+ # get query proj
+ query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2)
+
+ is_updated = False
+ if past_key_values is not None:
+ if isinstance(past_key_values, EncoderDecoderCache):
+ is_updated = past_key_values.is_updated.get(self.layer_idx)
+ if is_cross_attention:
+ # after the first generated id, we can subsequently re-use all key/value_states from cache
+ curr_past_key_values = past_key_values.cross_attention_cache
+ else:
+ curr_past_key_values = past_key_values.self_attention_cache
+ else:
+ curr_past_key_values = past_key_values
+
+ current_states = key_value_states if is_cross_attention else hidden_states
+ if is_cross_attention and past_key_values is not None and is_updated:
+ # reuse k,v, cross_attentions
+ key_states = curr_past_key_values.layers[self.layer_idx].keys
+ value_states = curr_past_key_values.layers[self.layer_idx].values
+ else:
+ key_states = self.k_proj(current_states)
+ value_states = self.v_proj(current_states)
+ key_states = key_states.view(*kv_input_shape).transpose(1, 2)
+ value_states = value_states.view(*kv_input_shape).transpose(1, 2)
+
+ if past_key_values is not None:
+ key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx)
+ # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+ if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):
+ past_key_values.is_updated[self.layer_idx] = True
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+class BartEncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: BartConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.embed_dim = config.d_model
+
+ self.self_attn = BartAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.encoder_attention_heads,
+ dropout=config.attention_dropout,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.dropout = config.dropout
+ self.activation_fn = ACT2FN[config.activation_function]
+ self.activation_dropout = config.activation_dropout
+ self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
+ self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor,
+ attention_mask: torch.FloatTensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states, _ = self.self_attn(
+ hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ residual = hidden_states
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+
+ if hidden_states.dtype == torch.float16 and not torch.isfinite(hidden_states).all():
+ clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+ hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+ return hidden_states
+
+
+class BartDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: BartConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.embed_dim = config.d_model
+
+ self.self_attn = BartAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.decoder_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ is_causal=True,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.dropout = config.dropout
+ self.activation_fn = ACT2FN[config.activation_function]
+ self.activation_dropout = config.activation_dropout
+
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.encoder_attn = BartAttention(
+ self.embed_dim,
+ config.decoder_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
+ self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = True,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ # Cross-Attention Block
+ if encoder_hidden_states is not None:
+ residual = hidden_states
+
+ hidden_states, _ = self.encoder_attn(
+ hidden_states,
+ key_value_states=encoder_hidden_states,
+ attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states = self.encoder_attn_layer_norm(hidden_states)
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+
+ return hidden_states
+
+
+class BartClassificationHead(nn.Module):
+ """Head for sentence-level classification tasks."""
+
+ def __init__(
+ self,
+ input_dim: int,
+ inner_dim: int,
+ num_classes: int,
+ pooler_dropout: float,
+ ):
+ super().__init__()
+ self.dense = nn.Linear(input_dim, inner_dim)
+ self.dropout = nn.Dropout(p=pooler_dropout)
+ self.out_proj = nn.Linear(inner_dim, num_classes)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.dense(hidden_states)
+ hidden_states = torch.tanh(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.out_proj(hidden_states)
+ return hidden_states
+
+
+@auto_docstring
+class BartPreTrainedModel(PreTrainedModel):
+ config: BartConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _keys_to_ignore_on_load_unexpected = ["encoder.version", "decoder.version"]
+ _no_split_modules = [r"BartEncoderLayer", r"BartDecoderLayer"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+
+ _can_compile_fullgraph = True
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, BartForConditionalGeneration):
+ init.zeros_(module.final_logits_bias)
+
+ @property
+ def dummy_inputs(self):
+ pad_token = self.config.pad_token_id
+ input_ids = torch.tensor([[0, 6, 10, 4, 2], [0, 8, 12, 2, pad_token]], device=self.device)
+ dummy_inputs = {
+ "attention_mask": input_ids.ne(pad_token),
+ "input_ids": input_ids,
+ }
+ return dummy_inputs
+
+
+class PretrainedBartModel(BartPreTrainedModel):
+ def __init_subclass__(self):
+ warnings.warn(
+ "The class `PretrainedBartModel` has been depreciated, please use `BartPreTrainedModel` instead.",
+ FutureWarning,
+ )
+
+
+class BartPretrainedModel(BartPreTrainedModel):
+ def __init_subclass__(self):
+ warnings.warn(
+ "The class `PretrainedBartModel` has been depreciated, please use `BartPreTrainedModel` instead.",
+ FutureWarning,
+ )
+
+
+class BartEncoder(BartPreTrainedModel):
+ """
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
+ [`BartEncoderLayer`].
+
+ Args:
+ config: BartConfig
+ embed_tokens (nn.Embedding): output embedding
+ """
+
+ _can_record_outputs = {
+ "hidden_states": BartEncoderLayer,
+ "attentions": BartAttention,
+ }
+
+ def __init__(self, config: BartConfig):
+ super().__init__(config)
+
+ self.dropout = config.dropout
+ self.layerdrop = config.encoder_layerdrop
+
+ embed_dim = config.d_model
+ self.padding_idx = config.pad_token_id
+ self.max_source_positions = config.max_position_embeddings
+ embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+
+ self.embed_tokens = BartScaledWordEmbedding(
+ config.vocab_size, embed_dim, self.padding_idx, embed_scale=embed_scale
+ )
+
+ self.embed_positions = BartLearnedPositionalEmbedding(
+ config.max_position_embeddings,
+ embed_dim,
+ )
+ self.layers = nn.ModuleList([BartEncoderLayer(config, layer_idx=i) for i in range(config.encoder_layers)])
+ self.layernorm_embedding = nn.LayerNorm(embed_dim)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutput:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ embed_pos = self.embed_positions(inputs_embeds[:, :, -1]) # needed for the shape only
+ embed_pos = embed_pos.to(inputs_embeds.device)
+
+ hidden_states = inputs_embeds + embed_pos
+ hidden_states = self.layernorm_embedding(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ )
+ for idx, encoder_layer in enumerate(self.layers):
+ # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+ to_drop = False
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop: # skip the layer
+ to_drop = True
+
+ if not to_drop:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask,
+ **kwargs,
+ )
+
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+class BartDecoder(BartPreTrainedModel):
+ """
+ Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`BartDecoderLayer`]
+
+ Args:
+ config: BartConfig
+ embed_tokens (nn.Embedding): output embedding
+ """
+
+ _can_record_outputs = {
+ "hidden_states": BartDecoderLayer,
+ "attentions": OutputRecorder(BartAttention, index=1, layer_name="self_attn"),
+ "cross_attentions": OutputRecorder(BartAttention, index=1, layer_name="encoder_attn"),
+ }
+
+ def __init__(self, config: BartConfig):
+ super().__init__(config)
+ self.dropout = config.dropout
+ self.layerdrop = config.decoder_layerdrop
+ self.padding_idx = config.pad_token_id
+ self.max_target_positions = config.max_position_embeddings
+ embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
+
+ self.embed_tokens = BartScaledWordEmbedding(
+ config.vocab_size, config.d_model, self.padding_idx, embed_scale=embed_scale
+ )
+
+ self.embed_positions = BartLearnedPositionalEmbedding(
+ config.max_position_embeddings,
+ config.d_model,
+ )
+ self.layers = nn.ModuleList([BartDecoderLayer(config, layer_idx=i) for i in range(config.decoder_layers)])
+
+ self.layernorm_embedding = nn.LayerNorm(config.d_model)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPastAndCrossAttentions:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of decoder_input_ids or decoder_inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ # initialize `past_key_values`
+ if use_cache and past_key_values is None:
+ past_key_values = (
+ EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+ if encoder_hidden_states is not None or self.config.is_encoder_decoder
+ else DynamicCache(config=self.config)
+ )
+
+ batch_size, seq_length = inputs_embeds.size()[:-1]
+ past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(seq_length, device=inputs_embeds.device) + past_key_values_length
+
+ if attention_mask is None and not is_torchdynamo_compiling():
+ # required mask seq length can be calculated via length of past cache
+ mask_seq_length = past_key_values_length + seq_length
+ attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)
+
+ self_attn_cache = (
+ past_key_values.self_attention_cache
+ if isinstance(past_key_values, EncoderDecoderCache)
+ else past_key_values
+ )
+
+ attention_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=self_attn_cache,
+ )
+ encoder_attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=encoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ )
+
+ # embed positions
+ positions = self.embed_positions(input, past_key_values_length, position_ids=position_ids)
+ positions = positions.to(inputs_embeds.device)
+
+ hidden_states = inputs_embeds + positions
+ hidden_states = self.layernorm_embedding(hidden_states)
+
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ for idx, decoder_layer in enumerate(self.layers):
+ # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop:
+ continue
+
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask,
+ encoder_hidden_states, # as a positional argument for gradient checkpointing
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class BartModel(BartPreTrainedModel):
+ _tied_weights_keys = {
+ "decoder.embed_tokens.weight": "shared.weight",
+ "encoder.embed_tokens.weight": "shared.weight",
+ }
+
+ def __init__(self, config: BartConfig):
+ super().__init__(config)
+
+ padding_idx, vocab_size = config.pad_token_id, config.vocab_size
+ embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
+ self.shared = BartScaledWordEmbedding(vocab_size, config.d_model, padding_idx, embed_scale=embed_scale)
+
+ self.encoder = BartEncoder(config)
+ self.decoder = BartDecoder(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.shared
+
+ def set_input_embeddings(self, value):
+ self.shared = value
+ self.encoder.embed_tokens = self.shared
+ self.decoder.embed_tokens = self.shared
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: list[torch.FloatTensor] | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ decoder_inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | Seq2SeqModelOutput:
+ r"""
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Indices of decoder input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+ Bart uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values`
+ is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`).
+
+ For translation and summarization training, `decoder_input_ids` should be provided. If no
+ `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to the right
+ for denoising pre-training following the paper.
+ decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+ be used by default.
+
+ If you want to change padding behavior, you should read [`modeling_bart._prepare_decoder_attention_mask`]
+ and modify to your needs. See diagram 1 in [the paper](https://huggingface.co/papers/1910.13461) for more
+ information on the default strategy.
+ """
+ # different to other models, Bart automatically creates decoder_input_ids from
+ # input_ids if no decoder_input_ids are provided
+ if decoder_input_ids is None and decoder_inputs_embeds is None:
+ if input_ids is None:
+ raise ValueError(
+ "If no `decoder_input_ids` or `decoder_inputs_embeds` are "
+ "passed, `input_ids` cannot be `None`. Please pass either "
+ "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`."
+ )
+
+ decoder_input_ids = shift_tokens_right(
+ input_ids, self.config.pad_token_id, self.config.decoder_start_token_id
+ )
+
+ if encoder_outputs is None:
+ encoder_outputs: BaseModelOutput = self.encoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+ elif not isinstance(encoder_outputs, BaseModelOutput):
+ encoder_outputs = BaseModelOutput(
+ last_hidden_state=encoder_outputs[0],
+ hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
+ attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
+ )
+
+ decoder_outputs: BaseModelOutputWithPastAndCrossAttentions = self.decoder(
+ input_ids=decoder_input_ids,
+ attention_mask=decoder_attention_mask,
+ encoder_hidden_states=encoder_outputs[0],
+ encoder_attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=decoder_inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ return Seq2SeqModelOutput(
+ last_hidden_state=decoder_outputs.last_hidden_state,
+ past_key_values=decoder_outputs.past_key_values,
+ decoder_hidden_states=decoder_outputs.hidden_states,
+ decoder_attentions=decoder_outputs.attentions,
+ cross_attentions=decoder_outputs.cross_attentions,
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+ encoder_hidden_states=encoder_outputs.hidden_states,
+ encoder_attentions=encoder_outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The BART Model with a language modeling head. Can be used for summarization.
+ """
+)
+class BartForConditionalGeneration(BartPreTrainedModel, GenerationMixin):
+ base_model_prefix = "model"
+ _tied_weights_keys = {
+ "lm_head.weight": "model.shared.weight",
+ }
+ _keys_to_ignore_on_load_missing = ["final_logits_bias"]
+
+ def __init__(self, config: BartConfig):
+ super().__init__(config)
+ self.model = BartModel(config)
+ self.register_buffer("final_logits_bias", torch.zeros((1, self.model.shared.num_embeddings)))
+ self.lm_head = nn.Linear(config.d_model, self.model.shared.num_embeddings, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def resize_token_embeddings(
+ self, new_num_tokens: int, pad_to_multiple_of: int | None = None, mean_resizing: bool = True
+ ) -> nn.Embedding:
+ new_embeddings = super().resize_token_embeddings(new_num_tokens, pad_to_multiple_of, mean_resizing)
+ self._resize_final_logits_bias(new_embeddings.weight.shape[0])
+ return new_embeddings
+
+ def _resize_final_logits_bias(self, new_num_tokens: int) -> None:
+ old_num_tokens = self.final_logits_bias.shape[-1]
+ if new_num_tokens <= old_num_tokens:
+ new_bias = self.final_logits_bias[:, :new_num_tokens]
+ else:
+ extra_bias = torch.zeros((1, new_num_tokens - old_num_tokens), device=self.final_logits_bias.device)
+ new_bias = torch.cat([self.final_logits_bias, extra_bias], dim=1)
+ self.register_buffer("final_logits_bias", new_bias)
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: list[torch.FloatTensor] | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ decoder_inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | Seq2SeqLMOutput:
+ r"""
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Indices of decoder input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+ Bart uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values`
+ is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`).
+
+ For translation and summarization training, `decoder_input_ids` should be provided. If no
+ `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to the right
+ for denoising pre-training following the paper.
+ decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+ be used by default.
+
+ If you want to change padding behavior, you should read [`modeling_bart._prepare_decoder_attention_mask`]
+ and modify to your needs. See diagram 1 in [the paper](https://huggingface.co/papers/1910.13461) for more
+ information on the default strategy.
+ 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]`.
+
+ Example summarization:
+
+ ```python
+ >>> from transformers import AutoTokenizer, BartForConditionalGeneration
+
+ >>> model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-cnn")
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
+
+ >>> ARTICLE_TO_SUMMARIZE = (
+ ... "PG&E stated it scheduled the blackouts in response to forecasts for high winds "
+ ... "amid dry conditions. The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were "
+ ... "scheduled to be affected by the shutoffs which were expected to last through at least midday tomorrow."
+ ... )
+ >>> inputs = tokenizer([ARTICLE_TO_SUMMARIZE], max_length=1024, return_tensors="pt")
+
+ >>> # Generate Summary
+ >>> summary_ids = model.generate(inputs["input_ids"], num_beams=2, min_length=0, max_length=20)
+ >>> tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ 'PG&E scheduled the blackouts in response to forecasts for high winds amid dry conditions'
+ ```
+
+ Mask filling example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, BartForConditionalGeneration
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
+ >>> model = BartForConditionalGeneration.from_pretrained("facebook/bart-base")
+
+ >>> TXT = "My friends are but they eat too many carbs."
+ >>> input_ids = tokenizer([TXT], return_tensors="pt")["input_ids"]
+ >>> logits = model(input_ids).logits
+
+ >>> masked_index = (input_ids[0] == tokenizer.mask_token_id).nonzero().item()
+ >>> probs = logits[0, masked_index].softmax(dim=0)
+ >>> values, predictions = probs.topk(5)
+
+ >>> tokenizer.decode(predictions).split()
+ ['not', 'good', 'healthy', 'great', 'very']
+ ```
+ """
+ if labels is not None:
+ if use_cache:
+ logger.warning("The `use_cache` argument is changed to `False` since `labels` is provided.")
+ use_cache = False
+ if decoder_input_ids is None and decoder_inputs_embeds is None:
+ decoder_input_ids = shift_tokens_right(
+ labels, self.config.pad_token_id, self.config.decoder_start_token_id
+ )
+
+ outputs: Seq2SeqModelOutput = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ decoder_input_ids=decoder_input_ids,
+ encoder_outputs=encoder_outputs,
+ decoder_attention_mask=decoder_attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ decoder_inputs_embeds=decoder_inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ lm_logits = self.lm_head(outputs[0])
+ lm_logits = lm_logits + self.final_logits_bias.to(lm_logits.device)
+
+ masked_lm_loss = None
+ if labels is not None:
+ labels = labels.to(lm_logits.device)
+ loss_fct = CrossEntropyLoss()
+ masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+ return Seq2SeqLMOutput(
+ loss=masked_lm_loss,
+ logits=lm_logits,
+ past_key_values=outputs.past_key_values,
+ decoder_hidden_states=outputs.decoder_hidden_states,
+ decoder_attentions=outputs.decoder_attentions,
+ cross_attentions=outputs.cross_attentions,
+ encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+ encoder_hidden_states=outputs.encoder_hidden_states,
+ encoder_attentions=outputs.encoder_attentions,
+ )
+
+ def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
+ return shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)
+
+
+@auto_docstring(
+ custom_intro="""
+ Bart model with a sequence classification/head on top (a linear layer on top of the pooled output) e.g. for GLUE
+ tasks.
+ """
+)
+class BartForSequenceClassification(BartPreTrainedModel):
+ def __init__(self, config: BartConfig, **kwargs):
+ super().__init__(config, **kwargs)
+ self.model = BartModel(config)
+ self.classification_head = BartClassificationHead(
+ config.d_model,
+ config.d_model,
+ config.num_labels,
+ config.classifier_dropout,
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: list[torch.FloatTensor] | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ decoder_inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | Seq2SeqSequenceClassifierOutput:
+ r"""
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Indices of decoder input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+ Bart uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values`
+ is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`).
+
+ For translation and summarization training, `decoder_input_ids` should be provided. If no
+ `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to the right
+ for denoising pre-training following the paper.
+ decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+ be used by default.
+
+ If you want to change padding behavior, you should read [`modeling_bart._prepare_decoder_attention_mask`]
+ and modify to your needs. See diagram 1 in [the paper](https://huggingface.co/papers/1910.13461) for more
+ information on the default strategy.
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ if labels is not None:
+ use_cache = False
+
+ if input_ids is None and inputs_embeds is not None:
+ raise NotImplementedError(
+ f"Passing input embeddings is currently not supported for {self.__class__.__name__}"
+ )
+
+ outputs: Seq2SeqModelOutput = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ decoder_input_ids=decoder_input_ids,
+ decoder_attention_mask=decoder_attention_mask,
+ encoder_outputs=encoder_outputs,
+ inputs_embeds=inputs_embeds,
+ decoder_inputs_embeds=decoder_inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+ hidden_states = outputs[0] # last hidden state
+
+ eos_mask = input_ids.eq(self.config.eos_token_id).to(hidden_states.device)
+
+ torch_compilable_check(
+ torch.unique_consecutive(eos_mask.sum(1)).numel() == 1,
+ "All examples must have the same number of tokens.",
+ )
+ sentence_representation = hidden_states[eos_mask, :].view(hidden_states.size(0), -1, hidden_states.size(-1))[
+ :, -1, :
+ ]
+ logits = self.classification_head(sentence_representation)
+
+ loss = None
+ if labels is not None:
+ labels = labels.to(logits.device)
+ if self.config.problem_type is None:
+ if self.config.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.config.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.config.num_labels == 1:
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.config.num_labels), labels.view(-1))
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(logits, labels)
+
+ return Seq2SeqSequenceClassifierOutput(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ decoder_hidden_states=outputs.decoder_hidden_states,
+ decoder_attentions=outputs.decoder_attentions,
+ cross_attentions=outputs.cross_attentions,
+ encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+ encoder_hidden_states=outputs.encoder_hidden_states,
+ encoder_attentions=outputs.encoder_attentions,
+ )
+
+
+@auto_docstring
+class BartForQuestionAnswering(BartPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ config.num_labels = 2
+ self.num_labels = config.num_labels
+
+ self.model = BartModel(config)
+ self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: list[torch.FloatTensor] | None = None,
+ start_positions: torch.LongTensor | None = None,
+ end_positions: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ decoder_inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | Seq2SeqQuestionAnsweringModelOutput:
+ r"""
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Indices of decoder input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+ Bart uses the `eos_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values`
+ is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`).
+
+ For translation and summarization training, `decoder_input_ids` should be provided. If no
+ `decoder_input_ids` is provided, the model will create this tensor by shifting the `input_ids` to the right
+ for denoising pre-training following the paper.
+ decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+ be used by default.
+
+ If you want to change padding behavior, you should read [`modeling_bart._prepare_decoder_attention_mask`]
+ and modify to your needs. See diagram 1 in [the paper](https://huggingface.co/papers/1910.13461) for more
+ information on the default strategy.
+ """
+ if start_positions is not None and end_positions is not None:
+ use_cache = False
+
+ outputs: Seq2SeqModelOutput = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ decoder_input_ids=decoder_input_ids,
+ decoder_attention_mask=decoder_attention_mask,
+ encoder_outputs=encoder_outputs,
+ inputs_embeds=inputs_embeds,
+ decoder_inputs_embeds=decoder_inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ sequence_output = outputs[0]
+
+ logits = self.qa_outputs(sequence_output)
+ start_logits, end_logits = logits.split(1, dim=-1)
+ start_logits = start_logits.squeeze(-1).contiguous()
+ end_logits = end_logits.squeeze(-1).contiguous()
+
+ total_loss = None
+ if start_positions is not None and end_positions is not None:
+ # If we are on multi-GPU, split add a dimension
+ if len(start_positions.size()) > 1:
+ start_positions = start_positions.squeeze(-1)
+ if len(end_positions.size()) > 1:
+ end_positions = end_positions.squeeze(-1)
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
+ ignored_index = start_logits.size(1)
+ start_positions = start_positions.clamp(0, ignored_index)
+ end_positions = end_positions.clamp(0, ignored_index)
+
+ loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+ start_loss = loss_fct(start_logits, start_positions)
+ end_loss = loss_fct(end_logits, end_positions)
+ total_loss = (start_loss + end_loss) / 2
+
+ return Seq2SeqQuestionAnsweringModelOutput(
+ loss=total_loss,
+ start_logits=start_logits,
+ end_logits=end_logits,
+ past_key_values=outputs.past_key_values,
+ decoder_hidden_states=outputs.decoder_hidden_states,
+ decoder_attentions=outputs.decoder_attentions,
+ cross_attentions=outputs.cross_attentions,
+ encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+ encoder_hidden_states=outputs.encoder_hidden_states,
+ encoder_attentions=outputs.encoder_attentions,
+ )
+
+
+class BartDecoderWrapper(BartPreTrainedModel):
+ """
+ This wrapper class is a helper class to correctly load pretrained checkpoints when the causal language model is
+ used in combination with the [`EncoderDecoderModel`] framework.
+ """
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.decoder = BartDecoder(config)
+ self.post_init()
+
+ def forward(self, *args, **kwargs):
+ return self.decoder(*args, **kwargs)
+
+
+@auto_docstring(
+ custom_intro="""
+ BART decoder with a language modeling head on top (linear layer with weights tied to the input embeddings).
+ """
+)
+class BartForCausalLM(BartPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {
+ "lm_head.weight": "model.decoder.embed_tokens.weight",
+ }
+
+ def __init__(self, config):
+ config.is_decoder = True
+ config.is_encoder_decoder = False
+ super().__init__(config)
+ self.model = BartDecoderWrapper(config)
+
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.decoder.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.model.decoder.embed_tokens = value
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithCrossAttentions:
+ r"""
+ 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]`.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, BartForCausalLM
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/bart-base")
+ >>> model = BartForCausalLM.from_pretrained("facebook/bart-base")
+ >>> assert model.config.is_decoder, f"{model.__class__} has to be configured as a decoder."
+ >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+ >>> outputs = model(**inputs)
+
+ >>> logits = outputs.logits
+ >>> expected_shape = [1, inputs.input_ids.shape[-1], model.config.vocab_size]
+ >>> list(logits.shape) == expected_shape
+ True
+ ```"""
+
+ outputs: BaseModelOutputWithPastAndCrossAttentions = self.model.decoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ # Only compute necessary logits
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ labels = labels.to(logits.device)
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+ return CausalLMOutputWithCrossAttentions(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+__all__ = [
+ "BartForCausalLM",
+ "BartForConditionalGeneration",
+ "BartForQuestionAnswering",
+ "BartForSequenceClassification",
+ "BartModel",
+ "BartPreTrainedModel",
+ "BartPretrainedModel",
+ "PretrainedBartModel",
+]
diff --git a/third_party/transformers/src/transformers/models/bart/tokenization_bart.py b/third_party/transformers/src/transformers/models/bart/tokenization_bart.py
new file mode 100644
index 0000000000000000000000000000000000000000..208d4e39131e87d7fb1f078f0693563d9c3d4631
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bart/tokenization_bart.py
@@ -0,0 +1,23 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# This source code is licensed under the Apache 2.0 license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""
+Compatibility shims for BART tokenizers in v5.
+
+In v5 we consolidate on the tokenizers-library backend and remove separate
+"slow" vs "fast" implementations. BART uses the same byte-level BPE
+tokenizer as RoBERTa, so we expose `BartTokenizer` and `BartTokenizerFast`
+as aliases to `RobertaTokenizer` to preserve the public API expected by
+existing code and tests.
+"""
+
+from ..roberta.tokenization_roberta import RobertaTokenizer as _RobertaTokenizer
+
+
+# Public aliases maintained for backwards compatibility
+BartTokenizer = _RobertaTokenizer
+BartTokenizerFast = _RobertaTokenizer
+
+__all__ = ["BartTokenizer", "BartTokenizerFast"]
diff --git a/third_party/transformers/src/transformers/models/bert/__init__.py b/third_party/transformers/src/transformers/models/bert/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..54a46312ad9153f2d8c3bf1a3029a3254efab791
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bert/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_bert import *
+ from .modeling_bert import *
+ from .tokenization_bert import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/bert/configuration_bert.py b/third_party/transformers/src/transformers/models/bert/configuration_bert.py
new file mode 100644
index 0000000000000000000000000000000000000000..1e3f288a820d0001fe5b2f7b3684f91a9b03cd2e
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bert/configuration_bert.py
@@ -0,0 +1,66 @@
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+"""BERT model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="google-bert/bert-base-uncased")
+@strict
+class BertConfig(PreTrainedConfig):
+ r"""
+ Examples:
+
+ ```python
+ >>> from transformers import BertConfig, BertModel
+
+ >>> # Initializing a BERT google-bert/bert-base-uncased style configuration
+ >>> configuration = BertConfig()
+
+ >>> # Initializing a model (with random weights) from the google-bert/bert-base-uncased style configuration
+ >>> model = BertModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "bert"
+
+ vocab_size: int = 30522
+ hidden_size: int = 768
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 12
+ intermediate_size: int = 3072
+ hidden_act: str = "gelu"
+ hidden_dropout_prob: float | int = 0.1
+ attention_probs_dropout_prob: float | int = 0.1
+ max_position_embeddings: int = 512
+ type_vocab_size: int = 2
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-12
+ pad_token_id: int | None = 0
+ use_cache: bool = True
+ classifier_dropout: float | int | None = None
+ is_decoder: bool = False
+ add_cross_attention: bool = False
+ bos_token_id: int | None = None
+ eos_token_id: int | list[int] | None = None
+ tie_word_embeddings: bool = True
+
+
+__all__ = ["BertConfig"]
diff --git a/third_party/transformers/src/transformers/models/bert/convert_bert_original_tf2_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/bert/convert_bert_original_tf2_checkpoint_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..9dfd8da474e37419472fe9b3ccd35d9836c671a1
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bert/convert_bert_original_tf2_checkpoint_to_pytorch.py
@@ -0,0 +1,246 @@
+# Copyright 2020 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+"""
+This script can be used to convert a head-less TF2.x Bert model to PyTorch, as published on the official (now
+deprecated) GitHub: https://github.com/tensorflow/models/tree/v2.3.0/official/nlp/bert
+
+TF2.x uses different variable names from the original BERT (TF 1.4) implementation. The script re-maps the TF2.x Bert
+weight names to the original names, so the model can be imported with Huggingface/transformer.
+
+You may adapt this script to include classification/MLM/NSP/etc. heads.
+
+Note: This script is only working with an older version of the TensorFlow models repository (<= v2.3.0).
+ Models trained with never versions are not compatible with this script.
+"""
+
+import argparse
+import os
+import re
+
+import tensorflow as tf
+import torch
+
+from transformers import BertConfig, BertModel
+from transformers.utils import logging
+
+
+logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+
+def load_tf2_weights_in_bert(model, tf_checkpoint_path, config):
+ tf_path = os.path.abspath(tf_checkpoint_path)
+ logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
+ # Load weights from TF model
+ init_vars = tf.train.list_variables(tf_path)
+ names = []
+ arrays = []
+ layer_depth = []
+ for full_name, shape in init_vars:
+ # logger.info(f"Loading TF weight {name} with shape {shape}")
+ name = full_name.split("/")
+ if full_name == "_CHECKPOINTABLE_OBJECT_GRAPH" or name[0] in ["global_step", "save_counter"]:
+ logger.info(f"Skipping non-model layer {full_name}")
+ continue
+ if "optimizer" in full_name:
+ logger.info(f"Skipping optimization layer {full_name}")
+ continue
+ if name[0] == "model":
+ # ignore initial 'model'
+ name = name[1:]
+ # figure out how many levels deep the name is
+ depth = 0
+ for _name in name:
+ if _name.startswith("layer_with_weights"):
+ depth += 1
+ else:
+ break
+ layer_depth.append(depth)
+ # read data
+ array = tf.train.load_variable(tf_path, full_name)
+ names.append("/".join(name))
+ arrays.append(array)
+ logger.info(f"Read a total of {len(arrays):,} layers")
+
+ # Sanity check
+ if len(set(layer_depth)) != 1:
+ raise ValueError(f"Found layer names with different depths (layer depth {list(set(layer_depth))})")
+ layer_depth = list(set(layer_depth))[0]
+ if layer_depth != 1:
+ raise ValueError(
+ "The model contains more than just the embedding/encoder layers. This script does not handle MLM/NSP"
+ " heads."
+ )
+
+ # convert layers
+ logger.info("Converting weights...")
+ for full_name, array in zip(names, arrays):
+ name = full_name.split("/")
+ pointer = model
+ trace = []
+ for i, m_name in enumerate(name):
+ if m_name == ".ATTRIBUTES":
+ # variable names end with .ATTRIBUTES/VARIABLE_VALUE
+ break
+ if m_name.startswith("layer_with_weights"):
+ layer_num = int(m_name.split("-")[-1])
+ if layer_num <= 2:
+ # embedding layers
+ # layer_num 0: word_embeddings
+ # layer_num 1: position_embeddings
+ # layer_num 2: token_type_embeddings
+ continue
+ elif layer_num == 3:
+ # embedding LayerNorm
+ trace.extend(["embeddings", "LayerNorm"])
+ pointer = getattr(pointer, "embeddings")
+ pointer = getattr(pointer, "LayerNorm")
+ elif layer_num > 3 and layer_num < config.num_hidden_layers + 4:
+ # encoder layers
+ trace.extend(["encoder", "layer", str(layer_num - 4)])
+ pointer = getattr(pointer, "encoder")
+ pointer = getattr(pointer, "layer")
+ pointer = pointer[layer_num - 4]
+ elif layer_num == config.num_hidden_layers + 4:
+ # pooler layer
+ trace.extend(["pooler", "dense"])
+ pointer = getattr(pointer, "pooler")
+ pointer = getattr(pointer, "dense")
+ elif m_name == "embeddings":
+ trace.append("embeddings")
+ pointer = getattr(pointer, "embeddings")
+ if layer_num == 0:
+ trace.append("word_embeddings")
+ pointer = getattr(pointer, "word_embeddings")
+ elif layer_num == 1:
+ trace.append("position_embeddings")
+ pointer = getattr(pointer, "position_embeddings")
+ elif layer_num == 2:
+ trace.append("token_type_embeddings")
+ pointer = getattr(pointer, "token_type_embeddings")
+ else:
+ raise ValueError(f"Unknown embedding layer with name {full_name}")
+ trace.append("weight")
+ pointer = getattr(pointer, "weight")
+ elif m_name == "_attention_layer":
+ # self-attention layer
+ trace.extend(["attention", "self"])
+ pointer = getattr(pointer, "attention")
+ pointer = getattr(pointer, "self")
+ elif m_name == "_attention_layer_norm":
+ # output attention norm
+ trace.extend(["attention", "output", "LayerNorm"])
+ pointer = getattr(pointer, "attention")
+ pointer = getattr(pointer, "output")
+ pointer = getattr(pointer, "LayerNorm")
+ elif m_name == "_attention_output_dense":
+ # output attention dense
+ trace.extend(["attention", "output", "dense"])
+ pointer = getattr(pointer, "attention")
+ pointer = getattr(pointer, "output")
+ pointer = getattr(pointer, "dense")
+ elif m_name == "_output_dense":
+ # output dense
+ trace.extend(["output", "dense"])
+ pointer = getattr(pointer, "output")
+ pointer = getattr(pointer, "dense")
+ elif m_name == "_output_layer_norm":
+ # output dense
+ trace.extend(["output", "LayerNorm"])
+ pointer = getattr(pointer, "output")
+ pointer = getattr(pointer, "LayerNorm")
+ elif m_name == "_key_dense":
+ # attention key
+ trace.append("key")
+ pointer = getattr(pointer, "key")
+ elif m_name == "_query_dense":
+ # attention query
+ trace.append("query")
+ pointer = getattr(pointer, "query")
+ elif m_name == "_value_dense":
+ # attention value
+ trace.append("value")
+ pointer = getattr(pointer, "value")
+ elif m_name == "_intermediate_dense":
+ # attention intermediate dense
+ trace.extend(["intermediate", "dense"])
+ pointer = getattr(pointer, "intermediate")
+ pointer = getattr(pointer, "dense")
+ elif m_name == "_output_layer_norm":
+ # output layer norm
+ trace.append("output")
+ pointer = getattr(pointer, "output")
+ # weights & biases
+ elif m_name in ["bias", "beta"]:
+ trace.append("bias")
+ pointer = getattr(pointer, "bias")
+ elif m_name in ["kernel", "gamma"]:
+ trace.append("weight")
+ pointer = getattr(pointer, "weight")
+ else:
+ logger.warning(f"Ignored {m_name}")
+ # for certain layers reshape is necessary
+ trace = ".".join(trace)
+ if re.match(r"(\S+)\.attention\.self\.(key|value|query)\.(bias|weight)", trace) or re.match(
+ r"(\S+)\.attention\.output\.dense\.weight", trace
+ ):
+ array = array.reshape(pointer.data.shape)
+ if "kernel" in full_name:
+ array = array.transpose()
+ if pointer.shape == array.shape:
+ pointer.data = torch.from_numpy(array)
+ else:
+ raise ValueError(
+ f"Shape mismatch in layer {full_name}: Model expects shape {pointer.shape} but layer contains shape:"
+ f" {array.shape}"
+ )
+ logger.info(f"Successfully set variable {full_name} to PyTorch layer {trace}")
+ return model
+
+
+def convert_tf2_checkpoint_to_pytorch(tf_checkpoint_path, config_path, pytorch_dump_path):
+ # Instantiate model
+ logger.info(f"Loading model based on config from {config_path}...")
+ config = BertConfig.from_json_file(config_path)
+ model = BertModel(config)
+
+ # Load weights from checkpoint
+ logger.info(f"Loading weights from checkpoint {tf_checkpoint_path}...")
+ load_tf2_weights_in_bert(model, tf_checkpoint_path, config)
+
+ # Save pytorch-model
+ logger.info(f"Saving PyTorch model to {pytorch_dump_path}...")
+ torch.save(model.state_dict(), pytorch_dump_path)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--tf_checkpoint_path", type=str, required=True, help="Path to the TensorFlow 2.x checkpoint path."
+ )
+ parser.add_argument(
+ "--bert_config_file",
+ type=str,
+ required=True,
+ help="The config json file corresponding to the BERT model. This specifies the model architecture.",
+ )
+ parser.add_argument(
+ "--pytorch_dump_path",
+ type=str,
+ required=True,
+ help="Path to the output PyTorch model (must include filename).",
+ )
+ args = parser.parse_args()
+ convert_tf2_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)
diff --git a/third_party/transformers/src/transformers/models/bert/convert_bert_original_tf_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/bert/convert_bert_original_tf_checkpoint_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..67775234673d2264c3fe88133d7272edf24df969
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bert/convert_bert_original_tf_checkpoint_to_pytorch.py
@@ -0,0 +1,61 @@
+# Copyright 2018 The HuggingFace Inc. team.
+#
+# 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.
+"""Convert BERT checkpoint."""
+
+import argparse
+
+import torch
+
+from transformers import BertConfig, BertForPreTraining, load_tf_weights_in_bert
+from transformers.utils import logging
+
+
+logging.set_verbosity_info()
+
+
+def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, bert_config_file, pytorch_dump_path):
+ # Initialise PyTorch model
+ config = BertConfig.from_json_file(bert_config_file)
+ print(f"Building PyTorch model from configuration: {config}")
+ model = BertForPreTraining(config)
+
+ # Load weights from tf checkpoint
+ load_tf_weights_in_bert(model, config, tf_checkpoint_path)
+
+ # Save pytorch-model
+ print(f"Save PyTorch model to {pytorch_dump_path}")
+ torch.save(model.state_dict(), pytorch_dump_path)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ # Required parameters
+ parser.add_argument(
+ "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
+ )
+ parser.add_argument(
+ "--bert_config_file",
+ default=None,
+ type=str,
+ required=True,
+ help=(
+ "The config json file corresponding to the pre-trained BERT model. \n"
+ "This specifies the model architecture."
+ ),
+ )
+ parser.add_argument(
+ "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
+ )
+ args = parser.parse_args()
+ convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)
diff --git a/third_party/transformers/src/transformers/models/bert/convert_bert_token_dropping_original_tf2_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/bert/convert_bert_token_dropping_original_tf2_checkpoint_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..a7832a53d55d452ed3aa0eec570a05c09a8e98e5
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bert/convert_bert_token_dropping_original_tf2_checkpoint_to_pytorch.py
@@ -0,0 +1,188 @@
+# Copyright 2022 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+"""
+This script converts a lm-head checkpoint from the "Token Dropping" implementation into a PyTorch-compatible BERT
+model. The official implementation of "Token Dropping" can be found in the TensorFlow Models repository:
+
+https://github.com/tensorflow/models/tree/master/official/projects/token_dropping
+"""
+
+import argparse
+
+import tensorflow as tf
+import torch
+
+from transformers import BertConfig, BertForMaskedLM
+from transformers.models.bert.modeling_bert import (
+ BertIntermediate,
+ BertLayer,
+ BertOutput,
+ BertPooler,
+ BertSelfAttention,
+ BertSelfOutput,
+)
+from transformers.utils import logging
+
+
+logging.set_verbosity_info()
+
+
+def convert_checkpoint_to_pytorch(tf_checkpoint_path: str, config_path: str, pytorch_dump_path: str):
+ def get_masked_lm_array(name: str):
+ full_name = f"masked_lm/{name}/.ATTRIBUTES/VARIABLE_VALUE"
+ array = tf.train.load_variable(tf_checkpoint_path, full_name)
+
+ if "kernel" in name:
+ array = array.transpose()
+
+ return torch.from_numpy(array)
+
+ def get_encoder_array(name: str):
+ full_name = f"encoder/{name}/.ATTRIBUTES/VARIABLE_VALUE"
+ array = tf.train.load_variable(tf_checkpoint_path, full_name)
+
+ if "kernel" in name:
+ array = array.transpose()
+
+ return torch.from_numpy(array)
+
+ def get_encoder_layer_array(layer_index: int, name: str):
+ full_name = f"encoder/_transformer_layers/{layer_index}/{name}/.ATTRIBUTES/VARIABLE_VALUE"
+ array = tf.train.load_variable(tf_checkpoint_path, full_name)
+
+ if "kernel" in name:
+ array = array.transpose()
+
+ return torch.from_numpy(array)
+
+ def get_encoder_attention_layer_array(layer_index: int, name: str, original_shape):
+ full_name = f"encoder/_transformer_layers/{layer_index}/_attention_layer/{name}/.ATTRIBUTES/VARIABLE_VALUE"
+ array = tf.train.load_variable(tf_checkpoint_path, full_name)
+ array = array.reshape(original_shape)
+
+ if "kernel" in name:
+ array = array.transpose()
+
+ return torch.from_numpy(array)
+
+ print(f"Loading model based on config from {config_path}...")
+ config = BertConfig.from_json_file(config_path)
+ model = BertForMaskedLM(config)
+
+ # Layers
+ for layer_index in range(0, config.num_hidden_layers):
+ layer: BertLayer = model.bert.encoder.layer[layer_index]
+
+ # Self-attention
+ self_attn: BertSelfAttention = layer.attention.self
+
+ self_attn.query.weight.data = get_encoder_attention_layer_array(
+ layer_index, "_query_dense/kernel", self_attn.query.weight.data.shape
+ )
+ self_attn.query.bias.data = get_encoder_attention_layer_array(
+ layer_index, "_query_dense/bias", self_attn.query.bias.data.shape
+ )
+ self_attn.key.weight.data = get_encoder_attention_layer_array(
+ layer_index, "_key_dense/kernel", self_attn.key.weight.data.shape
+ )
+ self_attn.key.bias.data = get_encoder_attention_layer_array(
+ layer_index, "_key_dense/bias", self_attn.key.bias.data.shape
+ )
+ self_attn.value.weight.data = get_encoder_attention_layer_array(
+ layer_index, "_value_dense/kernel", self_attn.value.weight.data.shape
+ )
+ self_attn.value.bias.data = get_encoder_attention_layer_array(
+ layer_index, "_value_dense/bias", self_attn.value.bias.data.shape
+ )
+
+ # Self-attention Output
+ self_output: BertSelfOutput = layer.attention.output
+
+ self_output.dense.weight.data = get_encoder_attention_layer_array(
+ layer_index, "_output_dense/kernel", self_output.dense.weight.data.shape
+ )
+ self_output.dense.bias.data = get_encoder_attention_layer_array(
+ layer_index, "_output_dense/bias", self_output.dense.bias.data.shape
+ )
+
+ self_output.LayerNorm.weight.data = get_encoder_layer_array(layer_index, "_attention_layer_norm/gamma")
+ self_output.LayerNorm.bias.data = get_encoder_layer_array(layer_index, "_attention_layer_norm/beta")
+
+ # Intermediate
+ intermediate: BertIntermediate = layer.intermediate
+
+ intermediate.dense.weight.data = get_encoder_layer_array(layer_index, "_intermediate_dense/kernel")
+ intermediate.dense.bias.data = get_encoder_layer_array(layer_index, "_intermediate_dense/bias")
+
+ # Output
+ bert_output: BertOutput = layer.output
+
+ bert_output.dense.weight.data = get_encoder_layer_array(layer_index, "_output_dense/kernel")
+ bert_output.dense.bias.data = get_encoder_layer_array(layer_index, "_output_dense/bias")
+
+ bert_output.LayerNorm.weight.data = get_encoder_layer_array(layer_index, "_output_layer_norm/gamma")
+ bert_output.LayerNorm.bias.data = get_encoder_layer_array(layer_index, "_output_layer_norm/beta")
+
+ # Embeddings
+ model.bert.embeddings.position_embeddings.weight.data = get_encoder_array("_position_embedding_layer/embeddings")
+ model.bert.embeddings.token_type_embeddings.weight.data = get_encoder_array("_type_embedding_layer/embeddings")
+ model.bert.embeddings.LayerNorm.weight.data = get_encoder_array("_embedding_norm_layer/gamma")
+ model.bert.embeddings.LayerNorm.bias.data = get_encoder_array("_embedding_norm_layer/beta")
+
+ # LM Head
+ lm_head = model.cls.predictions.transform
+
+ lm_head.dense.weight.data = get_masked_lm_array("dense/kernel")
+ lm_head.dense.bias.data = get_masked_lm_array("dense/bias")
+
+ lm_head.LayerNorm.weight.data = get_masked_lm_array("layer_norm/gamma")
+ lm_head.LayerNorm.bias.data = get_masked_lm_array("layer_norm/beta")
+
+ model.bert.embeddings.word_embeddings.weight.data = get_masked_lm_array("embedding_table")
+
+ # Pooling
+ model.bert.pooler = BertPooler(config=config)
+ model.bert.pooler.dense.weight.data: BertPooler = get_encoder_array("_pooler_layer/kernel")
+ model.bert.pooler.dense.bias.data: BertPooler = get_encoder_array("_pooler_layer/bias")
+
+ # Export final model
+ model.save_pretrained(pytorch_dump_path)
+
+ # Integration test - should load without any errors ;)
+ new_model = BertForMaskedLM.from_pretrained(pytorch_dump_path)
+ print(new_model.eval())
+
+ print("Model conversion was done successfully!")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--tf_checkpoint_path", type=str, required=True, help="Path to the TensorFlow Token Dropping checkpoint path."
+ )
+ parser.add_argument(
+ "--bert_config_file",
+ type=str,
+ required=True,
+ help="The config json file corresponding to the BERT model. This specifies the model architecture.",
+ )
+ parser.add_argument(
+ "--pytorch_dump_path",
+ type=str,
+ required=True,
+ help="Path to the output PyTorch model.",
+ )
+ args = parser.parse_args()
+ convert_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)
diff --git a/third_party/transformers/src/transformers/models/bert/modeling_bert.py b/third_party/transformers/src/transformers/models/bert/modeling_bert.py
new file mode 100644
index 0000000000000000000000000000000000000000..cc4adf29ebace537fe0dc30237a33416c2ba40e4
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bert/modeling_bert.py
@@ -0,0 +1,1395 @@
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# 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 BERT model."""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_bidirectional_mask, create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutputWithPastAndCrossAttentions,
+ BaseModelOutputWithPoolingAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions,
+ MaskedLMOutput,
+ MultipleChoiceModelOutput,
+ NextSentencePredictorOutput,
+ QuestionAnsweringModelOutput,
+ SequenceClassifierOutput,
+ TokenClassifierOutput,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...pytorch_utils import apply_chunking_to_forward
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, logging
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_bert import BertConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class BertEmbeddings(nn.Module):
+ """Construct the embeddings from word, position and token_type embeddings."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+ self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+ self.register_buffer(
+ "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+ )
+ self.register_buffer(
+ "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
+ )
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values_length: int = 0,
+ ) -> torch.Tensor:
+ if input_ids is not None:
+ input_shape = input_ids.size()
+ else:
+ input_shape = inputs_embeds.size()[:-1]
+
+ batch_size, seq_length = input_shape
+
+ if position_ids is None:
+ position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
+
+ # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
+ # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
+ # issue #5664
+ if token_type_ids is None:
+ if hasattr(self, "token_type_ids"):
+ # NOTE: We assume either pos ids to have bsz == 1 (broadcastable) or bsz == effective bsz (input_shape[0])
+ buffered_token_type_ids = self.token_type_ids.expand(position_ids.shape[0], -1)
+ buffered_token_type_ids = torch.gather(buffered_token_type_ids, dim=1, index=position_ids)
+ token_type_ids = buffered_token_type_ids.expand(batch_size, seq_length)
+ else:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.word_embeddings(input_ids)
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
+ embeddings = inputs_embeds + token_type_embeddings
+
+ position_embeddings = self.position_embeddings(position_ids)
+ embeddings = embeddings + position_embeddings
+
+ embeddings = self.LayerNorm(embeddings)
+ embeddings = self.dropout(embeddings)
+ return embeddings
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float | None = None,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ if scaling is None:
+ scaling = query.size(-1) ** -0.5
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class BertSelfAttention(nn.Module):
+ def __init__(self, config, is_causal=False, layer_idx=None):
+ super().__init__()
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+ raise ValueError(
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+ f"heads ({config.num_attention_heads})"
+ )
+ self.config = config
+
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+ self.scaling = self.attention_head_size**-0.5
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ self.is_decoder = config.is_decoder
+ self.is_causal = is_causal
+ self.layer_idx = layer_idx
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.attention_head_size)
+
+ # get all proj
+ query_layer = self.query(hidden_states).view(*hidden_shape).transpose(1, 2)
+ key_layer = self.key(hidden_states).view(*hidden_shape).transpose(1, 2)
+ value_layer = self.value(hidden_states).view(*hidden_shape).transpose(1, 2)
+
+ if past_key_values is not None:
+ # decoder-only bert can have a simple dynamic cache for example
+ current_past_key_values = past_key_values
+ if isinstance(past_key_values, EncoderDecoderCache):
+ current_past_key_values = past_key_values.self_attention_cache
+
+ # save all key/value_layer to cache to be re-used for fast auto-regressive generation
+ key_layer, value_layer = current_past_key_values.update(key_layer, value_layer, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_layer,
+ key_layer,
+ value_layer,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout.p,
+ scaling=self.scaling,
+ **kwargs,
+ )
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ return attn_output, attn_weights
+
+
+class BertCrossAttention(nn.Module):
+ def __init__(self, config, is_causal=False, layer_idx=None):
+ super().__init__()
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+ raise ValueError(
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+ f"heads ({config.num_attention_heads})"
+ )
+ self.config = config
+
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+ self.scaling = self.attention_head_size**-0.5
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ self.is_causal = is_causal
+ self.layer_idx = layer_idx
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ past_key_values: EncoderDecoderCache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor]:
+ # determine input shapes
+ bsz, tgt_len = hidden_states.shape[:-1]
+ src_len = encoder_hidden_states.shape[1]
+
+ q_input_shape = (bsz, tgt_len, -1, self.attention_head_size)
+ kv_input_shape = (bsz, src_len, -1, self.attention_head_size)
+
+ # get query proj
+ query_layer = self.query(hidden_states).view(*q_input_shape).transpose(1, 2)
+
+ is_updated = past_key_values.is_updated.get(self.layer_idx) if past_key_values is not None else False
+ if past_key_values is not None and is_updated:
+ # reuse k,v, cross_attentions
+ key_layer = past_key_values.cross_attention_cache.layers[self.layer_idx].keys
+ value_layer = past_key_values.cross_attention_cache.layers[self.layer_idx].values
+ else:
+ key_layer = self.key(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2)
+ value_layer = self.value(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2)
+
+ if past_key_values is not None:
+ # save all states to the cache
+ key_layer, value_layer = past_key_values.cross_attention_cache.update(
+ key_layer, value_layer, self.layer_idx
+ )
+ # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+ past_key_values.is_updated[self.layer_idx] = True
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_layer,
+ key_layer,
+ value_layer,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout.p,
+ scaling=self.scaling,
+ **kwargs,
+ )
+ attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()
+ return attn_output, attn_weights
+
+
+class BertSelfOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+class BertAttention(nn.Module):
+ def __init__(self, config, is_causal=False, layer_idx=None, is_cross_attention=False):
+ super().__init__()
+ self.is_cross_attention = is_cross_attention
+ attention_class = BertCrossAttention if is_cross_attention else BertSelfAttention
+ self.self = attention_class(config, is_causal=is_causal, layer_idx=layer_idx)
+ self.output = BertSelfOutput(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor]:
+ attention_mask = attention_mask if not self.is_cross_attention else encoder_attention_mask
+ attention_output, attn_weights = self.self(
+ hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ attention_output = self.output(attention_output, hidden_states)
+ return attention_output, attn_weights
+
+
+class BertIntermediate(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+class BertOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+class BertLayer(GradientCheckpointingLayer):
+ def __init__(self, config, layer_idx=None):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+ self.attention = BertAttention(config, is_causal=config.is_decoder, layer_idx=layer_idx)
+ self.is_decoder = config.is_decoder
+ self.add_cross_attention = config.add_cross_attention
+ if self.add_cross_attention:
+ if not self.is_decoder:
+ raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
+ self.crossattention = BertAttention(
+ config,
+ is_causal=False,
+ layer_idx=layer_idx,
+ is_cross_attention=True,
+ )
+ self.intermediate = BertIntermediate(config)
+ self.output = BertOutput(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ self_attention_output, _ = self.attention(
+ hidden_states,
+ attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ attention_output = self_attention_output
+
+ if self.is_decoder and encoder_hidden_states is not None:
+ if not hasattr(self, "crossattention"):
+ raise ValueError(
+ f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
+ " by setting `config.add_cross_attention=True`"
+ )
+
+ cross_attention_output, _ = self.crossattention(
+ self_attention_output,
+ None, # attention_mask
+ encoder_hidden_states,
+ encoder_attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ attention_output = cross_attention_output
+
+ layer_output = apply_chunking_to_forward(
+ self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+ )
+ return layer_output
+
+ def feed_forward_chunk(self, attention_output):
+ intermediate_output = self.intermediate(attention_output)
+ layer_output = self.output(intermediate_output, attention_output)
+ return layer_output
+
+
+class BertEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList([BertLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | BaseModelOutputWithPastAndCrossAttentions:
+ for i, layer_module in enumerate(self.layer):
+ hidden_states = layer_module(
+ hidden_states,
+ attention_mask,
+ encoder_hidden_states, # as a positional argument for gradient checkpointing
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values if use_cache else None,
+ )
+
+
+class BertPooler(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.activation = nn.Tanh()
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ # We "pool" the model by simply taking the hidden state corresponding
+ # to the first token.
+ first_token_tensor = hidden_states[:, 0]
+ pooled_output = self.dense(first_token_tensor)
+ pooled_output = self.activation(pooled_output)
+ return pooled_output
+
+
+class BertPredictionHeadTransform(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ if isinstance(config.hidden_act, str):
+ self.transform_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.transform_act_fn = config.hidden_act
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.transform_act_fn(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states)
+ return hidden_states
+
+
+class BertLMPredictionHead(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.transform = BertPredictionHeadTransform(config)
+
+ # The output weights are the same as the input embeddings, but there is
+ # an output-only bias for each token.
+ self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=True)
+ self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+ def forward(self, hidden_states):
+ hidden_states = self.transform(hidden_states)
+ hidden_states = self.decoder(hidden_states)
+ return hidden_states
+
+
+class BertOnlyMLMHead(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.predictions = BertLMPredictionHead(config)
+
+ def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
+ prediction_scores = self.predictions(sequence_output)
+ return prediction_scores
+
+
+class BertOnlyNSPHead(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.seq_relationship = nn.Linear(config.hidden_size, 2)
+
+ def forward(self, pooled_output):
+ seq_relationship_score = self.seq_relationship(pooled_output)
+ return seq_relationship_score
+
+
+class BertPreTrainingHeads(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.predictions = BertLMPredictionHead(config)
+ self.seq_relationship = nn.Linear(config.hidden_size, 2)
+
+ def forward(self, sequence_output, pooled_output):
+ prediction_scores = self.predictions(sequence_output)
+ seq_relationship_score = self.seq_relationship(pooled_output)
+ return prediction_scores, seq_relationship_score
+
+
+@auto_docstring
+class BertPreTrainedModel(PreTrainedModel):
+ config_class = BertConfig
+ base_model_prefix = "bert"
+ supports_gradient_checkpointing = True
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": BertLayer,
+ "attentions": BertSelfAttention,
+ "cross_attentions": BertCrossAttention,
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ super()._init_weights(module)
+ if isinstance(module, BertLMPredictionHead):
+ init.zeros_(module.bias)
+ elif isinstance(module, BertEmbeddings):
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+ init.zeros_(module.token_type_ids)
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Output type of [`BertForPreTraining`].
+ """
+)
+class BertForPreTrainingOutput(ModelOutput):
+ r"""
+ loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
+ Total loss as the sum of the masked language modeling loss and the next sequence prediction
+ (classification) loss.
+ 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).
+ seq_relationship_logits (`torch.FloatTensor` of shape `(batch_size, 2)`):
+ Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation
+ before SoftMax).
+ """
+
+ loss: torch.FloatTensor | None = None
+ prediction_logits: torch.FloatTensor | None = None
+ seq_relationship_logits: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+ cross-attention is added between the self-attention layers, following the architecture described in [Attention is
+ all you need](https://huggingface.co/papers/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
+ Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
+
+ To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set
+ to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and
+ `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.
+ """
+)
+class BertModel(BertPreTrainedModel):
+ _no_split_modules = ["BertEmbeddings", "BertLayer"]
+
+ def __init__(self, config, add_pooling_layer=True):
+ r"""
+ add_pooling_layer (bool, *optional*, defaults to `True`):
+ Whether to add a pooling layer
+ """
+ super().__init__(config)
+ self.config = config
+ self.gradient_checkpointing = False
+
+ self.embeddings = BertEmbeddings(config)
+ self.encoder = BertEncoder(config)
+
+ self.pooler = BertPooler(config) if add_pooling_layer else None
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embeddings.word_embeddings
+
+ def set_input_embeddings(self, value):
+ self.embeddings.word_embeddings = value
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | BaseModelOutputWithPoolingAndCrossAttentions:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if self.config.is_decoder:
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+ else:
+ use_cache = False
+
+ if use_cache and past_key_values is None:
+ past_key_values = (
+ EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+ if encoder_hidden_states is not None or self.config.is_encoder_decoder
+ else DynamicCache(config=self.config)
+ )
+
+ past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+
+ embedding_output = self.embeddings(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ past_key_values_length=past_key_values_length,
+ )
+
+ attention_mask, encoder_attention_mask = self._create_attention_masks(
+ attention_mask=attention_mask,
+ encoder_attention_mask=encoder_attention_mask,
+ embedding_output=embedding_output,
+ encoder_hidden_states=encoder_hidden_states,
+ past_key_values=past_key_values,
+ )
+
+ encoder_outputs = self.encoder(
+ embedding_output,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_ids=position_ids,
+ **kwargs,
+ )
+ sequence_output = encoder_outputs.last_hidden_state
+ pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+ return BaseModelOutputWithPoolingAndCrossAttentions(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ past_key_values=encoder_outputs.past_key_values,
+ )
+
+ def _create_attention_masks(
+ self,
+ attention_mask,
+ encoder_attention_mask,
+ embedding_output,
+ encoder_hidden_states,
+ past_key_values,
+ ):
+ if self.config.is_decoder:
+ attention_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=embedding_output,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ )
+ else:
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=embedding_output,
+ attention_mask=attention_mask,
+ )
+
+ if encoder_attention_mask is not None:
+ encoder_attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=embedding_output,
+ attention_mask=encoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ )
+
+ return attention_mask, encoder_attention_mask
+
+
+@auto_docstring(
+ custom_intro="""
+ Bert Model with two heads on top as done during the pretraining: a `masked language modeling` head and a `next
+ sentence prediction (classification)` head.
+ """
+)
+class BertForPreTraining(BertPreTrainedModel):
+ _tied_weights_keys = {
+ "cls.predictions.decoder.weight": "bert.embeddings.word_embeddings.weight",
+ "cls.predictions.decoder.bias": "cls.predictions.bias",
+ }
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.bert = BertModel(config)
+ self.cls = BertPreTrainingHeads(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.cls.predictions.decoder
+
+ def set_output_embeddings(self, new_embeddings):
+ self.cls.predictions.decoder = new_embeddings
+ self.cls.predictions.bias = new_embeddings.bias
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ next_sentence_label: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | BertForPreTrainingOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
+ config.vocab_size]` (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]`
+ next_sentence_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the next sequence prediction (classification) loss. Input should be a sequence
+ pair (see `input_ids` docstring) Indices should be in `[0, 1]`:
+
+ - 0 indicates sequence B is a continuation of sequence A,
+ - 1 indicates sequence B is a random sequence.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, BertForPreTraining
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
+ >>> model = BertForPreTraining.from_pretrained("google-bert/bert-base-uncased")
+
+ >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+ >>> outputs = model(**inputs)
+
+ >>> prediction_logits = outputs.prediction_logits
+ >>> seq_relationship_logits = outputs.seq_relationship_logits
+ ```
+ """
+ outputs = self.bert(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ return_dict=True,
+ **kwargs,
+ )
+
+ sequence_output, pooled_output = outputs[:2]
+ prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)
+
+ total_loss = None
+ if labels is not None and next_sentence_label is not None:
+ loss_fct = CrossEntropyLoss()
+ masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+ next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
+ total_loss = masked_lm_loss + next_sentence_loss
+
+ return BertForPreTrainingOutput(
+ loss=total_loss,
+ prediction_logits=prediction_scores,
+ seq_relationship_logits=seq_relationship_score,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Bert Model with a `language modeling` head on top for CLM fine-tuning.
+ """
+)
+class BertLMHeadModel(BertPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {
+ "cls.predictions.decoder.weight": "bert.embeddings.word_embeddings.weight",
+ "cls.predictions.decoder.bias": "cls.predictions.bias",
+ }
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ if not config.is_decoder:
+ logger.warning("If you want to use `BertLMHeadModel` as a standalone, add `is_decoder=True.`")
+
+ self.bert = BertModel(config, add_pooling_layer=False)
+ self.cls = BertOnlyMLMHead(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.cls.predictions.decoder
+
+ def set_output_embeddings(self, new_embeddings):
+ self.cls.predictions.decoder = new_embeddings
+ self.cls.predictions.bias = new_embeddings.bias
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | CausalLMOutputWithCrossAttentions:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
+ `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
+ ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`
+ """
+ if labels is not None:
+ use_cache = False
+
+ outputs: BaseModelOutputWithPoolingAndCrossAttentions = self.bert(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ return_dict=True,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.cls(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithCrossAttentions(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+@auto_docstring
+class BertForMaskedLM(BertPreTrainedModel):
+ _tied_weights_keys = {
+ "cls.predictions.decoder.weight": "bert.embeddings.word_embeddings.weight",
+ "cls.predictions.decoder.bias": "cls.predictions.bias",
+ }
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ if config.is_decoder:
+ logger.warning(
+ "If you want to use `BertForMaskedLM` make sure `config.is_decoder=False` for "
+ "bi-directional self-attention."
+ )
+
+ self.bert = BertModel(config, add_pooling_layer=False)
+ self.cls = BertOnlyMLMHead(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.cls.predictions.decoder
+
+ def set_output_embeddings(self, new_embeddings):
+ self.cls.predictions.decoder = new_embeddings
+ self.cls.predictions.bias = new_embeddings.bias
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | MaskedLMOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
+ config.vocab_size]` (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]`
+ """
+ outputs = self.bert(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ return_dict=True,
+ **kwargs,
+ )
+
+ sequence_output = outputs[0]
+ prediction_scores = self.cls(sequence_output)
+
+ masked_lm_loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss() # -100 index = padding token
+ masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+
+ return MaskedLMOutput(
+ loss=masked_lm_loss,
+ logits=prediction_scores,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Bert Model with a `next sentence prediction (classification)` head on top.
+ """
+)
+class BertForNextSentencePrediction(BertPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.bert = BertModel(config)
+ self.cls = BertOnlyNSPHead(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | NextSentencePredictorOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair
+ (see `input_ids` docstring). Indices should be in `[0, 1]`:
+
+ - 0 indicates sequence B is a continuation of sequence A,
+ - 1 indicates sequence B is a random sequence.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, BertForNextSentencePrediction
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
+ >>> model = BertForNextSentencePrediction.from_pretrained("google-bert/bert-base-uncased")
+
+ >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
+ >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light."
+ >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt")
+
+ >>> outputs = model(**encoding, labels=torch.LongTensor([1]))
+ >>> logits = outputs.logits
+ >>> assert logits[0, 0] < logits[0, 1] # next sentence was random
+ ```
+ """
+
+ outputs = self.bert(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ return_dict=True,
+ **kwargs,
+ )
+
+ pooled_output = outputs[1]
+
+ seq_relationship_scores = self.cls(pooled_output)
+
+ next_sentence_loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ next_sentence_loss = loss_fct(seq_relationship_scores.view(-1, 2), labels.view(-1))
+
+ return NextSentencePredictorOutput(
+ loss=next_sentence_loss,
+ logits=seq_relationship_scores,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Bert Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled
+ output) e.g. for GLUE tasks.
+ """
+)
+class BertForSequenceClassification(BertPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.config = config
+
+ self.bert = BertModel(config)
+ classifier_dropout = (
+ config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
+ )
+ self.dropout = nn.Dropout(classifier_dropout)
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | SequenceClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ outputs = self.bert(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ return_dict=True,
+ **kwargs,
+ )
+
+ pooled_output = outputs[1]
+
+ pooled_output = self.dropout(pooled_output)
+ logits = self.classifier(pooled_output)
+
+ loss = None
+ if labels is not None:
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(logits, labels)
+
+ return SequenceClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring
+class BertForMultipleChoice(BertPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.bert = BertModel(config)
+ classifier_dropout = (
+ config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
+ )
+ self.dropout = nn.Dropout(classifier_dropout)
+ self.classifier = nn.Linear(config.hidden_size, 1)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | MultipleChoiceModelOutput:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ token_type_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*):
+ Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
+ 1]`:
+
+ - 0 corresponds to a *sentence A* token,
+ - 1 corresponds to a *sentence B* token.
+
+ [What are token type IDs?](../glossary#token-type-ids)
+ position_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*):
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+ config.max_position_embeddings - 1]`.
+
+ [What are position IDs?](../glossary#position-ids)
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_choices, 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.
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
+ num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
+ `input_ids` above)
+ """
+ num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
+
+ input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
+ attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
+ token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
+ position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+ inputs_embeds = (
+ inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
+ if inputs_embeds is not None
+ else None
+ )
+
+ outputs = self.bert(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ return_dict=True,
+ **kwargs,
+ )
+
+ pooled_output = outputs[1]
+
+ pooled_output = self.dropout(pooled_output)
+ logits = self.classifier(pooled_output)
+ reshaped_logits = logits.view(-1, num_choices)
+
+ loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(reshaped_logits, labels)
+
+ return MultipleChoiceModelOutput(
+ loss=loss,
+ logits=reshaped_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring
+class BertForTokenClassification(BertPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.bert = BertModel(config, add_pooling_layer=False)
+ classifier_dropout = (
+ config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
+ )
+ self.dropout = nn.Dropout(classifier_dropout)
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | TokenClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
+ """
+ outputs = self.bert(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ return_dict=True,
+ **kwargs,
+ )
+
+ sequence_output = outputs[0]
+
+ sequence_output = self.dropout(sequence_output)
+ logits = self.classifier(sequence_output)
+
+ loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+ return TokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring
+class BertForQuestionAnswering(BertPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.bert = BertModel(config, add_pooling_layer=False)
+ self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ start_positions: torch.Tensor | None = None,
+ end_positions: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | QuestionAnsweringModelOutput:
+ outputs = self.bert(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ return_dict=True,
+ **kwargs,
+ )
+
+ sequence_output = outputs[0]
+
+ logits = self.qa_outputs(sequence_output)
+ start_logits, end_logits = logits.split(1, dim=-1)
+ start_logits = start_logits.squeeze(-1).contiguous()
+ end_logits = end_logits.squeeze(-1).contiguous()
+
+ total_loss = None
+ if start_positions is not None and end_positions is not None:
+ # If we are on multi-GPU, split add a dimension
+ if len(start_positions.size()) > 1:
+ start_positions = start_positions.squeeze(-1)
+ if len(end_positions.size()) > 1:
+ end_positions = end_positions.squeeze(-1)
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
+ ignored_index = start_logits.size(1)
+ start_positions = start_positions.clamp(0, ignored_index)
+ end_positions = end_positions.clamp(0, ignored_index)
+
+ loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+ start_loss = loss_fct(start_logits, start_positions)
+ end_loss = loss_fct(end_logits, end_positions)
+ total_loss = (start_loss + end_loss) / 2
+
+ return QuestionAnsweringModelOutput(
+ loss=total_loss,
+ start_logits=start_logits,
+ end_logits=end_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "BertForMaskedLM",
+ "BertForMultipleChoice",
+ "BertForNextSentencePrediction",
+ "BertForPreTraining",
+ "BertForQuestionAnswering",
+ "BertForSequenceClassification",
+ "BertForTokenClassification",
+ "BertLayer",
+ "BertLMHeadModel",
+ "BertModel",
+ "BertPreTrainedModel",
+]
diff --git a/third_party/transformers/src/transformers/models/bert/tokenization_bert.py b/third_party/transformers/src/transformers/models/bert/tokenization_bert.py
new file mode 100644
index 0000000000000000000000000000000000000000..47cb639fab40cfbfb9749230dc6b5c94d35d8b8d
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bert/tokenization_bert.py
@@ -0,0 +1,140 @@
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+#
+# 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.
+"""Tokenization classes for Bert."""
+
+import collections
+
+from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
+from tokenizers.models import WordPiece
+
+from ...tokenization_utils_tokenizers import TokenizersBackend
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt", "tokenizer_file": "tokenizer.json"}
+
+
+def load_vocab(vocab_file):
+ """Loads a vocabulary file into a dictionary."""
+ vocab = collections.OrderedDict()
+ with open(vocab_file, "r", encoding="utf-8") as reader:
+ tokens = reader.readlines()
+ for index, token in enumerate(tokens):
+ token = token.rstrip("\n")
+ vocab[token] = index
+ return vocab
+
+
+class BertTokenizer(TokenizersBackend):
+ r"""
+ Construct a BERT tokenizer (backed by HuggingFace's tokenizers library). Based on WordPiece.
+
+ This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should refer to
+ this superclass for more information regarding those methods.
+
+ Args:
+ vocab (`str` or `dict[str, int]`, *optional*):
+ Custom vocabulary dictionary. If not provided, vocabulary is loaded from `vocab_file`.
+ do_lower_case (`bool`, *optional*, defaults to `True`):
+ Whether or not to lowercase the input when tokenizing.
+ unk_token (`str`, *optional*, defaults to `"[UNK]"`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ sep_token (`str`, *optional*, defaults to `"[SEP]"`):
+ The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+ sequence classification or for a text and a question for question answering. It is also used as the last
+ token of a sequence built with special tokens.
+ pad_token (`str`, *optional*, defaults to `"[PAD]"`):
+ The token used for padding, for example when batching sequences of different lengths.
+ cls_token (`str`, *optional*, defaults to `"[CLS]"`):
+ The classifier token which is used when doing sequence classification (classification of the whole sequence
+ instead of per-token classification). It is the first token of the sequence when built with special tokens.
+ mask_token (`str`, *optional*, defaults to `"[MASK]"`):
+ The token used for masking values. This is the token used when training this model with masked language
+ modeling. This is the token which the model will try to predict.
+ tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
+ Whether or not to tokenize Chinese characters.
+ strip_accents (`bool`, *optional*):
+ Whether or not to strip all accents. If this option is not specified, then it will be determined by the
+ value for `lowercase` (as in the original BERT).
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ model_input_names = ["input_ids", "token_type_ids", "attention_mask"]
+ model = WordPiece
+
+ def __init__(
+ self,
+ vocab: str | dict[str, int] | None = None,
+ do_lower_case: bool = True,
+ unk_token: str = "[UNK]",
+ sep_token: str = "[SEP]",
+ pad_token: str = "[PAD]",
+ cls_token: str = "[CLS]",
+ mask_token: str = "[MASK]",
+ tokenize_chinese_chars: bool = True,
+ strip_accents: bool | None = None,
+ **kwargs,
+ ):
+ self.do_lower_case = do_lower_case
+ self.tokenize_chinese_chars = tokenize_chinese_chars
+ self.strip_accents = strip_accents
+ if vocab is None:
+ vocab = {
+ str(pad_token): 0,
+ str(unk_token): 1,
+ str(cls_token): 2,
+ str(sep_token): 3,
+ str(mask_token): 4,
+ }
+ self._vocab = vocab
+ self._tokenizer = Tokenizer(WordPiece(self._vocab, unk_token=str(unk_token)))
+ self._tokenizer.normalizer = normalizers.BertNormalizer(
+ clean_text=True,
+ handle_chinese_chars=tokenize_chinese_chars,
+ strip_accents=strip_accents,
+ lowercase=do_lower_case,
+ )
+ self._tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()
+ self._tokenizer.decoder = decoders.WordPiece(prefix="##")
+ super().__init__(
+ do_lower_case=do_lower_case,
+ unk_token=unk_token,
+ sep_token=sep_token,
+ pad_token=pad_token,
+ cls_token=cls_token,
+ mask_token=mask_token,
+ tokenize_chinese_chars=tokenize_chinese_chars,
+ strip_accents=strip_accents,
+ **kwargs,
+ )
+
+ cls_token_id = self.cls_token_id if self.cls_token_id is not None else 2
+ sep_token_id = self.sep_token_id if self.sep_token_id is not None else 3
+
+ self._tokenizer.post_processor = processors.TemplateProcessing(
+ single=f"{str(self.cls_token)}:0 $A:0 {str(self.sep_token)}:0",
+ pair=f"{str(self.cls_token)}:0 $A:0 {str(self.sep_token)}:0 $B:1 {str(self.sep_token)}:1",
+ special_tokens=[
+ (str(self.cls_token), cls_token_id),
+ (str(self.sep_token), sep_token_id),
+ ],
+ )
+
+
+__all__ = ["BertTokenizer"]
+
+BertTokenizerFast = BertTokenizer
diff --git a/third_party/transformers/src/transformers/models/bert/tokenization_bert_legacy.py b/third_party/transformers/src/transformers/models/bert/tokenization_bert_legacy.py
new file mode 100644
index 0000000000000000000000000000000000000000..cb96aa3eb29eee725edda716d00c8e56240c4e1e
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bert/tokenization_bert_legacy.py
@@ -0,0 +1,476 @@
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+#
+# 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.
+"""Tokenization classes for Bert."""
+
+import collections
+import os
+import unicodedata
+
+from ...tokenization_python import PreTrainedTokenizer, _is_control, _is_punctuation, _is_whitespace
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
+
+
+def load_vocab(vocab_file):
+ """Loads a vocabulary file into a dictionary."""
+ vocab = collections.OrderedDict()
+ with open(vocab_file, "r", encoding="utf-8") as reader:
+ tokens = reader.readlines()
+ for index, token in enumerate(tokens):
+ token = token.rstrip("\n")
+ vocab[token] = index
+ return vocab
+
+
+def whitespace_tokenize(text):
+ """Runs basic whitespace cleaning and splitting on a piece of text."""
+ text = text.strip()
+ if not text:
+ return []
+ tokens = text.split()
+ return tokens
+
+
+class BertTokenizerLegacy(PreTrainedTokenizer):
+ r"""
+ Construct a BERT tokenizer. Based on WordPiece.
+
+ This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+ this superclass for more information regarding those methods.
+
+ Args:
+ vocab_file (`str`):
+ File containing the vocabulary.
+ do_lower_case (`bool`, *optional*, defaults to `True`):
+ Whether or not to lowercase the input when tokenizing.
+ do_basic_tokenize (`bool`, *optional*, defaults to `True`):
+ Whether or not to do basic tokenization before WordPiece.
+ never_split (`Iterable`, *optional*):
+ Collection of tokens which will never be split during tokenization. Only has an effect when
+ `do_basic_tokenize=True`
+ unk_token (`str`, *optional*, defaults to `"[UNK]"`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ sep_token (`str`, *optional*, defaults to `"[SEP]"`):
+ The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+ sequence classification or for a text and a question for question answering. It is also used as the last
+ token of a sequence built with special tokens.
+ pad_token (`str`, *optional*, defaults to `"[PAD]"`):
+ The token used for padding, for example when batching sequences of different lengths.
+ cls_token (`str`, *optional*, defaults to `"[CLS]"`):
+ The classifier token which is used when doing sequence classification (classification of the whole sequence
+ instead of per-token classification). It is the first token of the sequence when built with special tokens.
+ mask_token (`str`, *optional*, defaults to `"[MASK]"`):
+ The token used for masking values. This is the token used when training this model with masked language
+ modeling. This is the token which the model will try to predict.
+ tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
+ Whether or not to tokenize Chinese characters.
+
+ This should likely be deactivated for Japanese (see this
+ [issue](https://github.com/huggingface/transformers/issues/328)).
+ strip_accents (`bool`, *optional*):
+ Whether or not to strip all accents. If this option is not specified, then it will be determined by the
+ value for `lowercase` (as in the original BERT).
+ clean_up_tokenization_spaces (`bool`, *optional*, defaults to `True`):
+ Whether or not to cleanup spaces after decoding, cleanup consists in removing potential artifacts like
+ extra spaces.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+
+ def __init__(
+ self,
+ vocab_file,
+ do_lower_case=True,
+ do_basic_tokenize=True,
+ never_split=None,
+ unk_token="[UNK]",
+ sep_token="[SEP]",
+ pad_token="[PAD]",
+ cls_token="[CLS]",
+ mask_token="[MASK]",
+ tokenize_chinese_chars=True,
+ strip_accents=None,
+ clean_up_tokenization_spaces=True,
+ **kwargs,
+ ):
+ if not os.path.isfile(vocab_file):
+ raise ValueError(
+ f"Can't find a vocabulary file at path '{vocab_file}'. To load the vocabulary from a Google pretrained"
+ " model use `tokenizer = BertTokenizer.from_pretrained(PRETRAINED_MODEL_NAME)`"
+ )
+ self.vocab = load_vocab(vocab_file)
+ self.ids_to_tokens = collections.OrderedDict([(ids, tok) for tok, ids in self.vocab.items()])
+ self.do_basic_tokenize = do_basic_tokenize
+ if do_basic_tokenize:
+ self.basic_tokenizer = BasicTokenizer(
+ do_lower_case=do_lower_case,
+ never_split=never_split,
+ tokenize_chinese_chars=tokenize_chinese_chars,
+ strip_accents=strip_accents,
+ )
+
+ self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab, unk_token=str(unk_token))
+
+ super().__init__(
+ do_lower_case=do_lower_case,
+ do_basic_tokenize=do_basic_tokenize,
+ never_split=never_split,
+ unk_token=unk_token,
+ sep_token=sep_token,
+ pad_token=pad_token,
+ cls_token=cls_token,
+ mask_token=mask_token,
+ tokenize_chinese_chars=tokenize_chinese_chars,
+ strip_accents=strip_accents,
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+ **kwargs,
+ )
+
+ @property
+ def do_lower_case(self):
+ return self.basic_tokenizer.do_lower_case
+
+ @property
+ def vocab_size(self):
+ return len(self.vocab)
+
+ def get_vocab(self):
+ return dict(self.vocab, **self.added_tokens_encoder)
+
+ def _tokenize(self, text, split_special_tokens=False):
+ split_tokens = []
+ if self.do_basic_tokenize:
+ for token in self.basic_tokenizer.tokenize(
+ text, never_split=self.all_special_tokens if not split_special_tokens else None
+ ):
+ # If the token is part of the never_split set
+ if token in self.basic_tokenizer.never_split:
+ split_tokens.append(token)
+ else:
+ split_tokens += self.wordpiece_tokenizer.tokenize(token)
+ else:
+ split_tokens = self.wordpiece_tokenizer.tokenize(text)
+ return split_tokens
+
+ def _convert_token_to_id(self, token):
+ """Converts a token (str) in an id using the vocab."""
+ return self.vocab.get(token, self.vocab.get(self.unk_token))
+
+ def _convert_id_to_token(self, index):
+ """Converts an index (integer) in a token (str) using the vocab."""
+ return self.ids_to_tokens.get(index, self.unk_token)
+
+ def convert_tokens_to_string(self, tokens):
+ """Converts a sequence of tokens (string) in a single string."""
+ out_string = " ".join(tokens).replace(" ##", "").strip()
+ return out_string
+
+ def build_inputs_with_special_tokens(
+ self, token_ids_0: list[int], token_ids_1: list[int] | None = None
+ ) -> list[int]:
+ """
+ Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+ adding special tokens. A BERT sequence has the following format:
+
+ - single sequence: `[CLS] X [SEP]`
+ - pair of sequences: `[CLS] A [SEP] B [SEP]`
+
+ Args:
+ token_ids_0 (`List[int]`):
+ List of IDs to which the special tokens will be added.
+ token_ids_1 (`List[int]`, *optional*):
+ Optional second list of IDs for sequence pairs.
+
+ Returns:
+ `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+ """
+ if token_ids_1 is None:
+ return [self.cls_token_id] + token_ids_0 + [self.sep_token_id]
+ cls = [self.cls_token_id]
+ sep = [self.sep_token_id]
+ return cls + token_ids_0 + sep + token_ids_1 + sep
+
+ def get_special_tokens_mask(
+ self, token_ids_0: list[int], token_ids_1: list[int] | None = None, already_has_special_tokens: bool = False
+ ) -> list[int]:
+ """
+ Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+ special tokens using the tokenizer `prepare_for_model` method.
+
+ Args:
+ token_ids_0 (`List[int]`):
+ List of IDs.
+ token_ids_1 (`List[int]`, *optional*):
+ Optional second list of IDs for sequence pairs.
+ already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+ Whether or not the token list is already formatted with special tokens for the model.
+
+ Returns:
+ `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+ """
+
+ if already_has_special_tokens:
+ return super().get_special_tokens_mask(
+ token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
+ )
+
+ if token_ids_1 is not None:
+ return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1)) + [1]
+ return [1] + ([0] * len(token_ids_0)) + [1]
+
+ def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
+ index = 0
+ if os.path.isdir(save_directory):
+ vocab_file = os.path.join(
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+ )
+ else:
+ vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
+ with open(vocab_file, "w", encoding="utf-8") as writer:
+ for token, token_index in sorted(self.vocab.items(), key=lambda kv: kv[1]):
+ if index != token_index:
+ logger.warning(
+ f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
+ " Please check that the vocabulary is not corrupted!"
+ )
+ index = token_index
+ writer.write(token + "\n")
+ index += 1
+ return (vocab_file,)
+
+
+class BasicTokenizer:
+ """
+ Constructs a BasicTokenizer that will run basic tokenization (punctuation splitting, lower casing, etc.).
+
+ Args:
+ do_lower_case (`bool`, *optional*, defaults to `True`):
+ Whether or not to lowercase the input when tokenizing.
+ never_split (`Iterable`, *optional*):
+ Collection of tokens which will never be split during tokenization. Only has an effect when
+ `do_basic_tokenize=True`
+ tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
+ Whether or not to tokenize Chinese characters.
+
+ This should likely be deactivated for Japanese (see this
+ [issue](https://github.com/huggingface/transformers/issues/328)).
+ strip_accents (`bool`, *optional*):
+ Whether or not to strip all accents. If this option is not specified, then it will be determined by the
+ value for `lowercase` (as in the original BERT).
+ do_split_on_punc (`bool`, *optional*, defaults to `True`):
+ In some instances we want to skip the basic punctuation splitting so that later tokenization can capture
+ the full context of the words, such as contractions.
+ """
+
+ def __init__(
+ self,
+ do_lower_case=True,
+ never_split=None,
+ tokenize_chinese_chars=True,
+ strip_accents=None,
+ do_split_on_punc=True,
+ ):
+ if never_split is None:
+ never_split = []
+ self.do_lower_case = do_lower_case
+ self.never_split = set(never_split)
+ self.tokenize_chinese_chars = tokenize_chinese_chars
+ self.strip_accents = strip_accents
+ self.do_split_on_punc = do_split_on_punc
+
+ def tokenize(self, text, never_split=None):
+ """
+ Basic Tokenization of a piece of text. For sub-word tokenization, see WordPieceTokenizer.
+
+ Args:
+ never_split (`List[str]`, *optional*)
+ Kept for backward compatibility purposes. Now implemented directly at the base class level (see
+ [`PreTrainedTokenizer.tokenize`]) List of token not to split.
+ """
+ # union() returns a new set by concatenating the two sets.
+ never_split = self.never_split.union(set(never_split)) if never_split else self.never_split
+ text = self._clean_text(text)
+
+ # This was added on November 1st, 2018 for the multilingual and Chinese
+ # models. This is also applied to the English models now, but it doesn't
+ # matter since the English models were not trained on any Chinese data
+ # and generally don't have any Chinese data in them (there are Chinese
+ # characters in the vocabulary because Wikipedia does have some Chinese
+ # words in the English Wikipedia.).
+ if self.tokenize_chinese_chars:
+ text = self._tokenize_chinese_chars(text)
+ # prevents treating the same character with different unicode codepoints as different characters
+ unicode_normalized_text = unicodedata.normalize("NFC", text)
+ orig_tokens = whitespace_tokenize(unicode_normalized_text)
+ split_tokens = []
+ for token in orig_tokens:
+ if token not in never_split:
+ if self.do_lower_case:
+ token = token.lower()
+ if self.strip_accents is not False:
+ token = self._run_strip_accents(token)
+ elif self.strip_accents:
+ token = self._run_strip_accents(token)
+ split_tokens.extend(self._run_split_on_punc(token, never_split))
+
+ output_tokens = whitespace_tokenize(" ".join(split_tokens))
+ return output_tokens
+
+ def _run_strip_accents(self, text):
+ """Strips accents from a piece of text."""
+ text = unicodedata.normalize("NFD", text)
+ output = []
+ for char in text:
+ cat = unicodedata.category(char)
+ if cat == "Mn":
+ continue
+ output.append(char)
+ return "".join(output)
+
+ def _run_split_on_punc(self, text, never_split=None):
+ """Splits punctuation on a piece of text."""
+ if not self.do_split_on_punc or (never_split is not None and text in never_split):
+ return [text]
+ chars = list(text)
+ i = 0
+ start_new_word = True
+ output = []
+ while i < len(chars):
+ char = chars[i]
+ if _is_punctuation(char):
+ output.append([char])
+ start_new_word = True
+ else:
+ if start_new_word:
+ output.append([])
+ start_new_word = False
+ output[-1].append(char)
+ i += 1
+
+ return ["".join(x) for x in output]
+
+ def _tokenize_chinese_chars(self, text):
+ """Adds whitespace around any CJK character."""
+ output = []
+ for char in text:
+ cp = ord(char)
+ if self._is_chinese_char(cp):
+ output.append(" ")
+ output.append(char)
+ output.append(" ")
+ else:
+ output.append(char)
+ return "".join(output)
+
+ def _is_chinese_char(self, cp):
+ """Checks whether CP is the codepoint of a CJK character."""
+ # This defines a "chinese character" as anything in the CJK Unicode block:
+ # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)
+ #
+ # Note that the CJK Unicode block is NOT all Japanese and Korean characters,
+ # despite its name. The modern Korean Hangul alphabet is a different block,
+ # as is Japanese Hiragana and Katakana. Those alphabets are used to write
+ # space-separated words, so they are not treated specially and handled
+ # like the all of the other languages.
+ if (
+ (cp >= 0x4E00 and cp <= 0x9FFF)
+ or (cp >= 0x3400 and cp <= 0x4DBF)
+ or (cp >= 0x20000 and cp <= 0x2A6DF)
+ or (cp >= 0x2A700 and cp <= 0x2B73F)
+ or (cp >= 0x2B740 and cp <= 0x2B81F)
+ or (cp >= 0x2B820 and cp <= 0x2CEAF)
+ or (cp >= 0xF900 and cp <= 0xFAFF)
+ or (cp >= 0x2F800 and cp <= 0x2FA1F)
+ ):
+ return True
+
+ return False
+
+ def _clean_text(self, text):
+ """Performs invalid character removal and whitespace cleanup on text."""
+ output = []
+ for char in text:
+ cp = ord(char)
+ if cp == 0 or cp == 0xFFFD or _is_control(char):
+ continue
+ if _is_whitespace(char):
+ output.append(" ")
+ else:
+ output.append(char)
+ return "".join(output)
+
+
+class WordpieceTokenizer:
+ """Runs WordPiece tokenization."""
+
+ def __init__(self, vocab, unk_token, max_input_chars_per_word=100):
+ self.vocab = vocab
+ self.unk_token = unk_token
+ self.max_input_chars_per_word = max_input_chars_per_word
+
+ def tokenize(self, text):
+ """
+ Tokenizes a piece of text into its word pieces. This uses a greedy longest-match-first algorithm to perform
+ tokenization using the given vocabulary.
+
+ For example, `input = "unaffable"` will return as output `["un", "##aff", "##able"]`.
+
+ Args:
+ text: A single token or whitespace separated tokens. This should have
+ already been passed through *BasicTokenizer*.
+
+ Returns:
+ A list of wordpiece tokens.
+ """
+
+ output_tokens = []
+ for token in whitespace_tokenize(text):
+ chars = list(token)
+ if len(chars) > self.max_input_chars_per_word:
+ output_tokens.append(self.unk_token)
+ continue
+
+ is_bad = False
+ start = 0
+ sub_tokens = []
+ while start < len(chars):
+ end = len(chars)
+ cur_substr = None
+ while start < end:
+ substr = "".join(chars[start:end])
+ if start > 0:
+ substr = "##" + substr
+ if substr in self.vocab:
+ cur_substr = substr
+ break
+ end -= 1
+ if cur_substr is None:
+ is_bad = True
+ break
+ sub_tokens.append(cur_substr)
+ start = end
+
+ if is_bad:
+ output_tokens.append(self.unk_token)
+ else:
+ output_tokens.extend(sub_tokens)
+ return output_tokens
+
+
+__all__ = ["BasicTokenizer", "BertTokenizerLegacy", "WordpieceTokenizer"]
diff --git a/third_party/transformers/src/transformers/models/bertweet/__init__.py b/third_party/transformers/src/transformers/models/bertweet/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..432622f1595d1a0d8bb1b3c9b9774b7d1e387d3e
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bertweet/__init__.py
@@ -0,0 +1,26 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .tokenization_bertweet import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/bertweet/tokenization_bertweet.py b/third_party/transformers/src/transformers/models/bertweet/tokenization_bertweet.py
new file mode 100644
index 0000000000000000000000000000000000000000..95e9f46787e8fb087e0bd1bac400f10a172e4710
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bertweet/tokenization_bertweet.py
@@ -0,0 +1,698 @@
+# Copyright (c) 2020, VinAI Research and the HuggingFace Inc. team.
+# Copyright 2018 The Open AI Team Authors and The HuggingFace Inc. team.
+#
+# 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.
+"""Tokenization classes for BERTweet"""
+
+import html
+import os
+import re
+
+import regex
+
+from ...tokenization_python import PreTrainedTokenizer
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {
+ "vocab_file": "vocab.txt",
+ "merges_file": "bpe.codes",
+}
+
+
+def get_pairs(word):
+ """
+ Return set of symbol pairs in a word.
+
+ Word is represented as tuple of symbols (symbols being variable-length strings).
+ """
+ pairs = set()
+ prev_char = word[0]
+ for char in word[1:]:
+ pairs.add((prev_char, char))
+ prev_char = char
+
+ pairs = set(pairs)
+ return pairs
+
+
+class BertweetTokenizer(PreTrainedTokenizer):
+ """
+ Constructs a BERTweet tokenizer, using Byte-Pair-Encoding.
+
+ This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+ this superclass for more information regarding those methods.
+
+ Args:
+ vocab_file (`str`):
+ Path to the vocabulary file.
+ merges_file (`str`):
+ Path to the merges file.
+ normalization (`bool`, *optional*, defaults to `False`):
+ Whether or not to apply a normalization preprocess.
+ bos_token (`str`, *optional*, defaults to `""`):
+ The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
+
+
+
+ When building a sequence using special tokens, this is not the token that is used for the beginning of
+ sequence. The token used is the `cls_token`.
+
+
+
+ eos_token (`str`, *optional*, defaults to `""`):
+ The end of sequence token.
+
+
+
+ When building a sequence using special tokens, this is not the token that is used for the end of sequence.
+ The token used is the `sep_token`.
+
+
+
+ sep_token (`str`, *optional*, defaults to `""`):
+ The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+ sequence classification or for a text and a question for question answering. It is also used as the last
+ token of a sequence built with special tokens.
+ cls_token (`str`, *optional*, defaults to `""`):
+ The classifier token which is used when doing sequence classification (classification of the whole sequence
+ instead of per-token classification). It is the first token of the sequence when built with special tokens.
+ unk_token (`str`, *optional*, defaults to `""`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ pad_token (`str`, *optional*, defaults to `""`):
+ The token used for padding, for example when batching sequences of different lengths.
+ mask_token (`str`, *optional*, defaults to `""`):
+ The token used for masking values. This is the token used when training this model with masked language
+ modeling. This is the token which the model will try to predict.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+
+ def __init__(
+ self,
+ vocab_file,
+ merges_file,
+ normalization=False,
+ bos_token="",
+ eos_token="",
+ sep_token="",
+ cls_token="",
+ unk_token="",
+ pad_token="",
+ mask_token="",
+ **kwargs,
+ ):
+ try:
+ from emoji import demojize
+
+ self.demojizer = demojize
+ except ImportError:
+ logger.warning(
+ "emoji is not installed, thus not converting emoticons or emojis into text. Install emoji: pip3"
+ " install emoji==0.6.0"
+ )
+ self.demojizer = None
+
+ self.vocab_file = vocab_file
+ self.merges_file = merges_file
+
+ self.encoder = {}
+ self.encoder[str(bos_token)] = 0
+ self.encoder[str(pad_token)] = 1
+ self.encoder[str(eos_token)] = 2
+ self.encoder[str(unk_token)] = 3
+
+ self.add_from_file(vocab_file)
+
+ self.decoder = {v: k for k, v in self.encoder.items()}
+
+ with open(merges_file, encoding="utf-8") as merges_handle:
+ merges = merges_handle.read().split("\n")[:-1]
+ merges = [tuple(merge.split()[:-1]) for merge in merges]
+ self.bpe_ranks = dict(zip(merges, range(len(merges))))
+ self.cache = {}
+
+ self.normalization = normalization
+ self.tweetPreprocessor = TweetTokenizer()
+ self.special_puncts = {"’": "'", "…": "..."}
+
+ super().__init__(
+ normalization=normalization,
+ bos_token=bos_token,
+ eos_token=eos_token,
+ sep_token=sep_token,
+ cls_token=cls_token,
+ unk_token=unk_token,
+ pad_token=pad_token,
+ mask_token=mask_token,
+ # Configure patterns instead of overriding methods
+ token_type_ids_pattern="all_zeros", # BERTweet doesn't use token type IDs
+ token_type_ids_include_special_tokens=True,
+ special_tokens_pattern="cls_double_sep", # X Y
+ **kwargs,
+ )
+
+ @property
+ def vocab_size(self):
+ return len(self.encoder)
+
+ def get_vocab(self):
+ return dict(self.encoder, **self.added_tokens_encoder)
+
+ def bpe(self, token):
+ if token in self.cache:
+ return self.cache[token]
+ word = tuple(token)
+ word = tuple(list(word[:-1]) + [word[-1] + ""])
+ pairs = get_pairs(word)
+
+ if not pairs:
+ return token
+
+ while True:
+ bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
+ if bigram not in self.bpe_ranks:
+ break
+ first, second = bigram
+ new_word = []
+ i = 0
+ while i < len(word):
+ try:
+ j = word.index(first, i)
+ except ValueError:
+ new_word.extend(word[i:])
+ break
+ else:
+ new_word.extend(word[i:j])
+ i = j
+
+ if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
+ new_word.append(first + second)
+ i += 2
+ else:
+ new_word.append(word[i])
+ i += 1
+ new_word = tuple(new_word)
+ word = new_word
+ if len(word) == 1:
+ break
+ else:
+ pairs = get_pairs(word)
+ word = "@@ ".join(word)
+ word = word[:-4]
+ self.cache[token] = word
+ return word
+
+ def _tokenize(self, text):
+ """Tokenize a string."""
+ if self.normalization: # Perform Tweet normalization before performing BPE
+ text = self.normalizeTweet(text)
+
+ split_tokens = []
+ words = re.findall(r"\S+\n?", text)
+ for token in words:
+ split_tokens.extend(list(self.bpe(token).split(" ")))
+ return split_tokens
+
+ def normalizeTweet(self, tweet):
+ """
+ Normalize a raw Tweet
+ """
+ for punct in self.special_puncts:
+ tweet = tweet.replace(punct, self.special_puncts[punct])
+
+ tokens = self.tweetPreprocessor.tokenize(tweet)
+ normTweet = " ".join([self.normalizeToken(token) for token in tokens])
+
+ normTweet = (
+ normTweet.replace("cannot ", "can not ")
+ .replace("n't ", " n't ")
+ .replace("n 't ", " n't ")
+ .replace("ca n't", "can't")
+ .replace("ai n't", "ain't")
+ )
+ normTweet = (
+ normTweet.replace("'m ", " 'm ")
+ .replace("'re ", " 're ")
+ .replace("'s ", " 's ")
+ .replace("'ll ", " 'll ")
+ .replace("'d ", " 'd ")
+ .replace("'ve ", " 've ")
+ )
+ normTweet = (
+ normTweet.replace(" p . m .", " p.m.")
+ .replace(" p . m ", " p.m ")
+ .replace(" a . m .", " a.m.")
+ .replace(" a . m ", " a.m ")
+ )
+
+ return " ".join(normTweet.split())
+
+ def normalizeToken(self, token):
+ """
+ Normalize tokens in a Tweet
+ """
+ lowercased_token = token.lower()
+ if token.startswith("@"):
+ return "@USER"
+ elif lowercased_token.startswith("http") or lowercased_token.startswith("www"):
+ return "HTTPURL"
+ elif len(token) == 1:
+ if token in self.special_puncts:
+ return self.special_puncts[token]
+ if self.demojizer is not None:
+ return self.demojizer(token)
+ else:
+ return token
+ else:
+ return token
+
+ def _convert_token_to_id(self, token):
+ """Converts a token (str) in an id using the vocab."""
+ return self.encoder.get(token, self.encoder.get(self.unk_token))
+
+ def _convert_id_to_token(self, index):
+ """Converts an index (integer) in a token (str) using the vocab."""
+ return self.decoder.get(index, self.unk_token)
+
+ def convert_tokens_to_string(self, tokens):
+ """Converts a sequence of tokens (string) in a single string."""
+ out_string = " ".join(tokens).replace("@@ ", "").strip()
+ return out_string
+
+ # def decode(self, token_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True):
+ # filtered_tokens = ' '.join(self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens))
+ # tokens_generated_so_far = re.sub('(@@ )', '', string=filtered_tokens)
+ # tokens_generated_so_far = re.sub('(@@ ?$)', '', string=tokens_generated_so_far)
+ # return ''.join(tokens_generated_so_far)
+
+ def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str, ...]:
+ """
+ Save the vocabulary and merges files to a directory.
+ """
+ if not os.path.isdir(save_directory):
+ logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+ return ()
+
+ vocab_files_names = getattr(self, "vocab_files_names", {})
+ prefix = f"{filename_prefix}-" if filename_prefix else ""
+
+ # Save vocabulary in the format expected by add_from_file:
+ # Exclude special tokens (IDs 0-3) as they are added in __init__ before add_from_file
+ vocab_file = os.path.join(save_directory, prefix + vocab_files_names.get("vocab_file", "vocab.txt"))
+ with open(vocab_file, "w", encoding="utf-8") as f:
+ for token, token_id in sorted(self.encoder.items(), key=lambda kv: kv[1]):
+ # Only save tokens with ID >= 4, as IDs 0-3 are reserved for special tokens
+ if token_id >= 4:
+ f.write(f"{token} {token_id}\n")
+
+ # Save BPE merges
+ merge_file = os.path.join(save_directory, prefix + vocab_files_names.get("merges_file", "bpe.codes"))
+ with open(merge_file, "w", encoding="utf-8") as writer:
+ writer.writelines(
+ " ".join(bpe_tokens) + "\n"
+ for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1])
+ )
+
+ return (vocab_file, merge_file)
+
+ def add_from_file(self, f):
+ """
+ Loads a pre-existing dictionary from a text file and adds its symbols to this instance.
+ """
+ if isinstance(f, str):
+ try:
+ with open(f, "r", encoding="utf-8") as fd:
+ self.add_from_file(fd)
+ except FileNotFoundError as fnfe:
+ raise fnfe
+ except UnicodeError:
+ raise Exception(f"Incorrect encoding detected in {f}, please rebuild the dataset")
+ return
+
+ lines = f.readlines()
+ for lineTmp in lines:
+ line = lineTmp.strip()
+ idx = line.rfind(" ")
+ if idx == -1:
+ raise ValueError("Incorrect dictionary format, expected ' '")
+ word = line[:idx]
+ self.encoder[word] = len(self.encoder)
+
+
+# Natural Language Toolkit: Twitter Tokenizer
+#
+# Copyright (C) 2001-2020 NLTK Project
+# Author: Christopher Potts
+# Ewan Klein (modifications)
+# Pierpaolo Pantone <> (modifications)
+# URL: http://nltk.org/
+# For license information, see LICENSE.TXT
+#
+
+
+"""
+Twitter-aware tokenizer, designed to be flexible and easy to adapt to new domains and tasks. The basic logic is this:
+
+1. The tuple regex_strings defines a list of regular expression strings.
+
+2. The regex_strings strings are put, in order, into a compiled regular expression object called word_re.
+
+3. The tokenization is done by word_re.findall(s), where s is the user-supplied string, inside the tokenize() method of
+ the class Tokenizer.
+
+4. When instantiating Tokenizer objects, there is a single option: preserve_case. By default, it is set to True. If it
+ is set to False, then the tokenizer will lowercase everything except for emoticons.
+
+"""
+
+
+######################################################################
+#
+# import regex # https://github.com/nltk/nltk/issues/2409
+# import html
+#
+######################################################################
+# The following strings are components in the regular expression
+# that is used for tokenizing. It's important that phone_number
+# appears first in the final regex (since it can contain whitespace).
+# It also could matter that tags comes after emoticons, due to the
+# possibility of having text like
+#
+# <:| and some text >:)
+#
+# Most importantly, the final element should always be last, since it
+# does a last ditch whitespace-based tokenization of whatever is left.
+
+# ToDo: Update with http://en.wikipedia.org/wiki/List_of_emoticons ?
+
+# This particular element is used in a couple ways, so we define it
+# with a name:
+# docstyle-ignore
+EMOTICONS = r"""
+ (?:
+ [<>]?
+ [:;=8] # eyes
+ [\-o\*\']? # optional nose
+ [\)\]\(\[dDpP/\:\}\{@\|\\] # mouth
+ |
+ [\)\]\(\[dDpP/\:\}\{@\|\\] # mouth
+ [\-o\*\']? # optional nose
+ [:;=8] # eyes
+ [<>]?
+ |
+ <3 # heart
+ )"""
+
+# URL pattern due to John Gruber, modified by Tom Winzig. See
+# https://gist.github.com/winzig/8894715
+# docstyle-ignore
+URLS = r""" # Capture 1: entire matched URL
+ (?:
+ https?: # URL protocol and colon
+ (?:
+ /{1,3} # 1-3 slashes
+ | # or
+ [a-z0-9%] # Single letter or digit or '%'
+ # (Trying not to match e.g. "URI::Escape")
+ )
+ | # or
+ # looks like domain name followed by a slash:
+ [a-z0-9.\-]+[.]
+ (?:[a-z]{2,13})
+ /
+ )
+ (?: # One or more:
+ [^\s()<>{}\[\]]+ # Run of non-space, non-()<>{}[]
+ | # or
+ \([^\s()]*?\([^\s()]+\)[^\s()]*?\) # balanced parens, one level deep: (...(...)...)
+ |
+ \([^\s]+?\) # balanced parens, non-recursive: (...)
+ )+
+ (?: # End with:
+ \([^\s()]*?\([^\s()]+\)[^\s()]*?\) # balanced parens, one level deep: (...(...)...)
+ |
+ \([^\s]+?\) # balanced parens, non-recursive: (...)
+ | # or
+ [^\s`!()\[\]{};:'".,<>?«»“”‘’] # not a space or one of these punct chars
+ )
+ | # OR, the following to match naked domains:
+ (?:
+ (?\s]+>""",
+ # ASCII Arrows
+ r"""[\-]+>|<[\-]+""",
+ # Twitter username:
+ r"""(?:@[\w_]+)""",
+ # Twitter hashtags:
+ r"""(?:\#+[\w_]+[\w\'_\-]*[\w_]+)""",
+ # email addresses
+ r"""[\w.+-]+@[\w-]+\.(?:[\w-]\.?)+[\w-]""",
+ # docstyle-ignore
+ # Remaining word types:
+ r"""
+ (?:[^\W\d_](?:[^\W\d_]|['\-_])+[^\W\d_]) # Words with apostrophes or dashes.
+ |
+ (?:[+\-]?\d+[,/.:-]\d+[+\-]?) # Numbers, including fractions, decimals.
+ |
+ (?:[\w_]+) # Words without apostrophes or dashes.
+ |
+ (?:\.(?:\s*\.){1,}) # Ellipsis dots.
+ |
+ (?:\S) # Everything else that isn't whitespace.
+ """,
+)
+
+######################################################################
+# This is the core tokenizing regex:
+
+WORD_RE = regex.compile(r"""(%s)""" % "|".join(REGEXPS), regex.VERBOSE | regex.I | regex.UNICODE)
+
+# WORD_RE performs poorly on these patterns:
+HANG_RE = regex.compile(r"([^a-zA-Z0-9])\1{3,}")
+
+# The emoticon string gets its own regex so that we can preserve case for
+# them as needed:
+EMOTICON_RE = regex.compile(EMOTICONS, regex.VERBOSE | regex.I | regex.UNICODE)
+
+# These are for regularizing HTML entities to Unicode:
+ENT_RE = regex.compile(r"&(#?(x?))([^&;\s]+);")
+
+
+######################################################################
+# Functions for converting html entities
+######################################################################
+
+
+def _str_to_unicode(text, encoding=None, errors="strict"):
+ if encoding is None:
+ encoding = "utf-8"
+ if isinstance(text, bytes):
+ return text.decode(encoding, errors)
+ return text
+
+
+def _replace_html_entities(text, keep=(), remove_illegal=True, encoding="utf-8"):
+ """
+ Remove entities from text by converting them to their corresponding unicode character.
+
+ Args:
+ text:
+ A unicode string or a byte string encoded in the given *encoding* (which defaults to 'utf-8').
+ keep (list):
+ List of entity names which should not be replaced. This supports both numeric entities (`&#nnnn;` and
+ `&#hhhh;`) and named entities (such as ` ` or `>`).
+ remove_illegal (bool):
+ If `True`, entities that can't be converted are removed. Otherwise, entities that can't be converted are
+ kept "as is".
+ Returns: A unicode string with the entities removed.
+
+ See https://github.com/scrapy/w3lib/blob/master/w3lib/html.py
+
+ Examples:
+
+ ```python
+ >>> from nltk.tokenize.casual import _replace_html_entities
+
+ >>> _replace_html_entities(b"Price: £100")
+ 'Price: \\xa3100'
+
+ >>> print(_replace_html_entities(b"Price: £100"))
+ Price: £100
+ ```"""
+
+ def _convert_entity(match):
+ entity_body = match.group(3)
+ if match.group(1):
+ try:
+ if match.group(2):
+ number = int(entity_body, 16)
+ else:
+ number = int(entity_body, 10)
+ # Numeric character references in the 80-9F range are typically
+ # interpreted by browsers as representing the characters mapped
+ # to bytes 80-9F in the Windows-1252 encoding. For more info
+ # see: https://en.wikipedia.org/wiki/ISO/IEC_8859-1#Similar_character_sets
+ if 0x80 <= number <= 0x9F:
+ return bytes((number,)).decode("cp1252")
+ except ValueError:
+ number = None
+ else:
+ if entity_body in keep:
+ return match.group(0)
+ else:
+ number = html.entities.name2codepoint.get(entity_body)
+ if number is not None:
+ try:
+ return chr(number)
+ except (ValueError, OverflowError):
+ pass
+
+ return "" if remove_illegal else match.group(0)
+
+ return ENT_RE.sub(_convert_entity, _str_to_unicode(text, encoding))
+
+
+######################################################################
+
+
+class TweetTokenizer:
+ r"""
+ Examples:
+
+ ```python
+ >>> # Tokenizer for tweets.
+ >>> from nltk.tokenize import TweetTokenizer
+
+ >>> tknzr = TweetTokenizer()
+ >>> s0 = "This is a cooool #dummysmiley: :-) :-P <3 and some arrows < > -> <--"
+ >>> tknzr.tokenize(s0)
+ ['This', 'is', 'a', 'cooool', '#dummysmiley', ':', ':-)', ':-P', '<3', 'and', 'some', 'arrows', '<', '>', '->', '<--']
+
+ >>> # Examples using *strip_handles* and *reduce_len parameters*:
+ >>> tknzr = TweetTokenizer(strip_handles=True, reduce_len=True)
+ >>> s1 = "@remy: This is waaaaayyyy too much for you!!!!!!"
+ >>> tknzr.tokenize(s1)
+ [':', 'This', 'is', 'waaayyy', 'too', 'much', 'for', 'you', '!', '!', '!']
+ ```"""
+
+ def __init__(self, preserve_case=True, reduce_len=False, strip_handles=False):
+ self.preserve_case = preserve_case
+ self.reduce_len = reduce_len
+ self.strip_handles = strip_handles
+
+ def tokenize(self, text):
+ """
+ Args:
+ text: str
+
+ Returns: list(str) A tokenized list of strings; concatenating this list returns the original string if
+ `preserve_case=False`
+ """
+ # Fix HTML character entities:
+ text = _replace_html_entities(text)
+ # Remove username handles
+ if self.strip_handles:
+ text = remove_handles(text)
+ # Normalize word lengthening
+ if self.reduce_len:
+ text = reduce_lengthening(text)
+ # Shorten problematic sequences of characters
+ safe_text = HANG_RE.sub(r"\1\1\1", text)
+ # Tokenize:
+ words = WORD_RE.findall(safe_text)
+ # Possibly alter the case, but avoid changing emoticons like :D into :d:
+ if not self.preserve_case:
+ words = [x if EMOTICON_RE.search(x) else x.lower() for x in words]
+ return words
+
+
+######################################################################
+# Normalization Functions
+######################################################################
+
+
+def reduce_lengthening(text):
+ """
+ Replace repeated character sequences of length 3 or greater with sequences of length 3.
+ """
+ pattern = regex.compile(r"(.)\1{2,}")
+ return pattern.sub(r"\1\1\1", text)
+
+
+def remove_handles(text):
+ """
+ Remove Twitter username handles from text.
+ """
+ pattern = regex.compile(
+ r"(?>> from transformers import BioGptModel, BioGptConfig
+
+ >>> # Initializing a BioGPT microsoft/biogpt style configuration
+ >>> configuration = BioGptConfig()
+
+ >>> # Initializing a model from the microsoft/biogpt style configuration
+ >>> model = BioGptModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "biogpt"
+
+ vocab_size: int = 42384
+ hidden_size: int = 1024
+ num_hidden_layers: int = 24
+ num_attention_heads: int = 16
+ intermediate_size: int = 4096
+ hidden_act: str = "gelu"
+ hidden_dropout_prob: float | int = 0.1
+ attention_probs_dropout_prob: float | int = 0.1
+ max_position_embeddings: int = 1024
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-12
+ scale_embedding: bool = True
+ use_cache: bool = True
+ layerdrop: float | int = 0.0
+ activation_dropout: float | int = 0.0
+ pad_token_id: int | None = 1
+ bos_token_id: int | None = 0
+ eos_token_id: int | list[int] | None = 2
+ tie_word_embeddings: bool = True
+
+
+__all__ = ["BioGptConfig"]
diff --git a/third_party/transformers/src/transformers/models/biogpt/convert_biogpt_original_pytorch_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/biogpt/convert_biogpt_original_pytorch_checkpoint_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..e8c82bba94a23014fcb7d62e99c84e3cbfd082ae
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/biogpt/convert_biogpt_original_pytorch_checkpoint_to_pytorch.py
@@ -0,0 +1,291 @@
+# Copyright 2022 The HuggingFace Inc. team.
+#
+# 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.
+
+
+import argparse
+import json
+import os
+import re
+import shutil
+
+import torch
+
+from transformers import BioGptConfig, BioGptForCausalLM
+from transformers.models.biogpt.tokenization_biogpt import VOCAB_FILES_NAMES
+from transformers.tokenization_utils_base import TOKENIZER_CONFIG_FILE
+from transformers.utils import WEIGHTS_NAME, logging
+
+
+logging.set_verbosity_warning()
+
+json_indent = 2
+
+
+# modified from https://github.com/facebookresearch/fairseq/blob/dd74992d0d143155998e9ed4076826bcea80fb06/fairseq/data/dictionary.py#L18
+class Dictionary:
+ """A mapping from symbols to consecutive integers"""
+
+ def __init__(
+ self,
+ *, # begin keyword-only arguments
+ bos="",
+ pad="",
+ eos="",
+ unk="",
+ extra_special_symbols=None,
+ ):
+ self.bos_word, self.unk_word, self.pad_word, self.eos_word = bos, unk, pad, eos
+ self.symbols = []
+ self.count = []
+ self.indices = {}
+ self.bos_index = self.add_symbol(bos)
+ self.pad_index = self.add_symbol(pad)
+ self.eos_index = self.add_symbol(eos)
+ self.unk_index = self.add_symbol(unk)
+ if extra_special_symbols:
+ for s in extra_special_symbols:
+ self.add_symbol(s)
+ self.nspecial = len(self.symbols)
+
+ def __eq__(self, other):
+ return self.indices == other.indices
+
+ def __getitem__(self, idx):
+ if idx < len(self.symbols):
+ return self.symbols[idx]
+ return self.unk_word
+
+ def __len__(self):
+ """Returns the number of symbols in the dictionary"""
+ return len(self.symbols)
+
+ def __contains__(self, sym):
+ return sym in self.indices
+
+ @classmethod
+ def load(cls, f):
+ """Loads the dictionary from a text file with the format:
+
+ ```
+
+
+ ...
+ ```
+ """
+ d = cls()
+ d.add_from_file(f)
+ return d
+
+ def add_symbol(self, word, n=1, overwrite=False):
+ """Adds a word to the dictionary"""
+ if word in self.indices and not overwrite:
+ idx = self.indices[word]
+ self.count[idx] = self.count[idx] + n
+ return idx
+ else:
+ idx = len(self.symbols)
+ self.indices[word] = idx
+ self.symbols.append(word)
+ self.count.append(n)
+ return idx
+
+ def _load_meta(self, lines):
+ return 0
+
+ def add_from_file(self, f):
+ """
+ Loads a pre-existing dictionary from a text file and adds its symbols to this instance.
+ """
+ if isinstance(f, str):
+ try:
+ with open(f, "r", encoding="utf-8") as fd:
+ self.add_from_file(fd)
+ except FileNotFoundError as fnfe:
+ raise fnfe
+ except UnicodeError:
+ raise Exception(f"Incorrect encoding detected in {f}, please rebuild the dataset")
+ return
+
+ lines = f.readlines()
+ indices_start_line = self._load_meta(lines)
+
+ for line in lines[indices_start_line:]:
+ try:
+ line, field = line.rstrip().rsplit(" ", 1)
+ if field == "#fairseq:overwrite":
+ overwrite = True
+ line, field = line.rsplit(" ", 1)
+ else:
+ overwrite = False
+ count = int(field)
+ word = line
+ if word in self and not overwrite:
+ raise RuntimeError(
+ f"Duplicate word found when loading Dictionary: '{word}'. "
+ "Duplicate words can overwrite earlier ones by adding the "
+ "#fairseq:overwrite flag at the end of the corresponding row "
+ "in the dictionary file. If using the Camembert model, please "
+ "download an updated copy of the model file."
+ )
+ self.add_symbol(word, n=count, overwrite=overwrite)
+ except ValueError:
+ raise ValueError("Incorrect dictionary format, expected ' [flags]'")
+
+
+def rewrite_dict_keys(d):
+ # (1) remove word breaking symbol, (2) add word ending symbol where the word is not broken up,
+ # e.g.: d = {'le@@': 5, 'tt@@': 6, 'er': 7} => {'le': 5, 'tt': 6, 'er': 7}
+ d2 = dict((re.sub(r"@@$", "", k), v) if k.endswith("@@") else (re.sub(r"$", "", k), v) for k, v in d.items())
+ keep_keys = ["", "", "", ""]
+ # restore the special tokens
+ for k in keep_keys:
+ del d2[f"{k}"]
+ d2[k] = d[k] # restore
+ return d2
+
+
+def convert_biogpt_checkpoint_to_pytorch(biogpt_checkpoint_path, pytorch_dump_folder_path):
+ # prep
+ if not os.path.exists(biogpt_checkpoint_path):
+ raise ValueError(f"path {biogpt_checkpoint_path} does not exist!")
+ os.makedirs(pytorch_dump_folder_path, exist_ok=True)
+ print(f"Writing results to {pytorch_dump_folder_path}")
+
+ # handle various types of models
+
+ checkpoint_file = os.path.join(biogpt_checkpoint_path, "checkpoint.pt")
+ if not os.path.isfile(checkpoint_file):
+ raise ValueError(f"path to the file {checkpoint_file} does not exist!")
+ chkpt = torch.load(checkpoint_file, map_location="cpu", weights_only=True)
+
+ args = chkpt["cfg"]["model"]
+
+ # dicts
+ dict_file = os.path.join(biogpt_checkpoint_path, "dict.txt")
+ if not os.path.isfile(dict_file):
+ raise ValueError(f"path to the file {dict_file} does not exist!")
+ src_dict = Dictionary.load(dict_file)
+ src_vocab = rewrite_dict_keys(src_dict.indices)
+ src_vocab_size = len(src_vocab)
+ src_vocab_file = os.path.join(pytorch_dump_folder_path, VOCAB_FILES_NAMES["vocab_file"])
+ print(f"Generating {src_vocab_file} of {src_vocab_size} records")
+ with open(src_vocab_file, "w", encoding="utf-8") as f:
+ f.write(json.dumps(src_vocab, ensure_ascii=False, indent=json_indent))
+
+ # merges_file (bpecodes)
+ bpecodes_file = os.path.join(biogpt_checkpoint_path, "bpecodes")
+ if not os.path.isfile(bpecodes_file):
+ raise ValueError(f"path to the file {bpecodes_file} does not exist!")
+
+ merges_file = os.path.join(pytorch_dump_folder_path, VOCAB_FILES_NAMES["merges_file"])
+ shutil.copyfile(bpecodes_file, merges_file)
+
+ # model config
+ biogpt_model_config_file = os.path.join(pytorch_dump_folder_path, "config.json")
+
+ model_conf = {
+ "activation_dropout": args["activation_dropout"],
+ "architectures": ["BioGptForCausalLM"],
+ "attention_probs_dropout_prob": args["attention_dropout"],
+ "bos_token_id": 0,
+ "eos_token_id": 2,
+ "hidden_act": args["activation_fn"],
+ "hidden_dropout_prob": args["dropout"],
+ "hidden_size": args["decoder_embed_dim"],
+ "initializer_range": 0.02,
+ "intermediate_size": args["decoder_ffn_embed_dim"],
+ "layer_norm_eps": 1e-12,
+ "layerdrop": args["decoder_layerdrop"],
+ "max_position_embeddings": args["max_target_positions"],
+ "model_type": "biogpt",
+ "num_attention_heads": args["decoder_attention_heads"],
+ "num_hidden_layers": args["decoder_layers"],
+ "pad_token_id": 1,
+ "scale_embedding": not args["no_scale_embedding"],
+ "tie_word_embeddings": args["share_decoder_input_output_embed"],
+ "vocab_size": src_vocab_size,
+ }
+
+ # good hparam defaults to start with
+
+ print(f"Generating {biogpt_model_config_file}")
+ with open(biogpt_model_config_file, "w", encoding="utf-8") as f:
+ f.write(json.dumps(model_conf, ensure_ascii=False, indent=json_indent))
+
+ # tokenizer config
+ biogpt_tokenizer_config_file = os.path.join(pytorch_dump_folder_path, TOKENIZER_CONFIG_FILE)
+
+ tokenizer_conf = {
+ "bos_token": "",
+ "eos_token": "",
+ "model_max_length": 1024,
+ "pad_token": "",
+ "special_tokens_map_file": None,
+ "tokenizer_class": "BioGptTokenizer",
+ "unk_token": "",
+ }
+
+ print(f"Generating {biogpt_tokenizer_config_file}")
+ with open(biogpt_tokenizer_config_file, "w", encoding="utf-8") as f:
+ f.write(json.dumps(tokenizer_conf, ensure_ascii=False, indent=json_indent))
+
+ # model
+ model_state_dict = chkpt["model"]
+
+ # remove unneeded keys
+ ignore_keys = [
+ "decoder.version",
+ ]
+ for k in ignore_keys:
+ model_state_dict.pop(k, None)
+
+ layer_names = list(model_state_dict.keys())
+ for layer_name in layer_names:
+ if layer_name.endswith("output_projection.weight"):
+ model_state_dict[layer_name.replace("decoder.", "")] = model_state_dict.pop(layer_name)
+ else:
+ model_state_dict[layer_name.replace("decoder", "biogpt")] = model_state_dict.pop(layer_name)
+
+ config = BioGptConfig.from_pretrained(pytorch_dump_folder_path)
+ model_new = BioGptForCausalLM(config)
+
+ # check that it loads ok
+ model_new.load_state_dict(model_state_dict)
+
+ # save
+ pytorch_weights_dump_path = os.path.join(pytorch_dump_folder_path, WEIGHTS_NAME)
+ print(f"Generating {pytorch_weights_dump_path}")
+ torch.save(model_state_dict, pytorch_weights_dump_path)
+
+ print("Conversion is done!")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ # Required parameters
+ parser.add_argument(
+ "--biogpt_checkpoint_path",
+ default=None,
+ type=str,
+ required=True,
+ help=(
+ "Path to the official PyTorch checkpoint file which is expected to reside in the dump dir with dicts,"
+ " bpecodes, etc."
+ ),
+ )
+ parser.add_argument(
+ "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
+ )
+ args = parser.parse_args()
+ convert_biogpt_checkpoint_to_pytorch(args.biogpt_checkpoint_path, args.pytorch_dump_folder_path)
diff --git a/third_party/transformers/src/transformers/models/biogpt/modeling_biogpt.py b/third_party/transformers/src/transformers/models/biogpt/modeling_biogpt.py
new file mode 100644
index 0000000000000000000000000000000000000000..e1a7e2e3fdaf4b837c62ba63a3033bc47d06e411
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/biogpt/modeling_biogpt.py
@@ -0,0 +1,678 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/biogpt/modular_biogpt.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_biogpt.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2022 The HuggingFace Team and Microsoft Research AI4Science All rights reserved.
+#
+# 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.
+
+import math
+from collections.abc import Callable
+
+import torch
+import torch.nn as nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutputWithPastAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions,
+ SequenceClassifierOutputWithPast,
+ TokenClassifierOutput,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_biogpt import BioGptConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class BioGptLearnedPositionalEmbedding(nn.Embedding):
+ """
+ This module learns positional embeddings up to a fixed maximum size.
+ """
+
+ def __init__(self, num_embeddings: int, embedding_dim: int):
+ # BIOGPT is set up so that if padding_idx is specified then offset the embedding ids by 2
+ # and adjust num_embeddings appropriately. Other models don't have this hack
+ self.offset = 2
+ super().__init__(num_embeddings + self.offset, embedding_dim)
+
+ def forward(
+ self,
+ attention_mask: torch.LongTensor,
+ past_key_values_length: int = 0,
+ position_ids: torch.LongTensor | None = None,
+ ):
+ """`input_ids_shape` is expected to be [bsz x seqlen]."""
+
+ if position_ids is None:
+ position_ids = torch.cumsum(attention_mask, dim=1)
+ position_ids = (position_ids * attention_mask - 1).long()
+ # cut positions if `past_key_values_length` is > 0
+ position_ids = position_ids[:, past_key_values_length:]
+
+ return super().forward(position_ids + self.offset)
+
+
+class BioGptScaledWordEmbedding(nn.Embedding):
+ """
+ This module overrides nn.Embeddings' forward by multiplying with embeddings scale.
+ """
+
+ def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: float | None = 1.0):
+ super().__init__(num_embeddings, embedding_dim, padding_idx)
+ self.embed_scale = embed_scale
+
+ def forward(self, input_ids: torch.Tensor):
+ return super().forward(input_ids) * self.embed_scale
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float | None = None,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ if scaling is None:
+ scaling = query.size(-1) ** -0.5
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class BioGptAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ dropout: float = 0.0,
+ is_decoder: bool = False,
+ bias: bool = True,
+ is_causal: bool = False,
+ config: BioGptConfig | None = None,
+ layer_idx: int | None = None,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.dropout = dropout
+ self.head_dim = embed_dim // num_heads
+ self.config = config
+
+ if (self.head_dim * num_heads) != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+ f" and `num_heads`: {num_heads})."
+ )
+ self.scaling = self.head_dim**-0.5
+ self.is_decoder = is_decoder
+ self.is_causal = is_causal
+ self.layer_idx = layer_idx
+ if layer_idx is None and self.is_decoder:
+ logger.warning_once(
+ f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and "
+ "will lead to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class."
+ )
+
+ self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ key_value_states: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ # TODO: we need a refactor so that the different attention modules can get their specific kwargs
+ # ATM, we have mixed things encoder, decoder, and encoder-decoder attn
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Input shape: Batch x Time x Channel"""
+
+ # if key_value_states are provided this layer is used as a cross-attention layer
+ # for the decoder
+ is_cross_attention = key_value_states is not None
+
+ # determine input shapes
+ bsz, tgt_len = hidden_states.shape[:-1]
+ src_len = key_value_states.shape[1] if is_cross_attention else tgt_len
+
+ q_input_shape = (bsz, tgt_len, -1, self.head_dim)
+ kv_input_shape = (bsz, src_len, -1, self.head_dim)
+
+ # get query proj
+ query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2)
+
+ is_updated = False
+ if past_key_values is not None:
+ if isinstance(past_key_values, EncoderDecoderCache):
+ is_updated = past_key_values.is_updated.get(self.layer_idx)
+ if is_cross_attention:
+ # after the first generated id, we can subsequently re-use all key/value_states from cache
+ curr_past_key_values = past_key_values.cross_attention_cache
+ else:
+ curr_past_key_values = past_key_values.self_attention_cache
+ else:
+ curr_past_key_values = past_key_values
+
+ current_states = key_value_states if is_cross_attention else hidden_states
+ if is_cross_attention and past_key_values is not None and is_updated:
+ # reuse k,v, cross_attentions
+ key_states = curr_past_key_values.layers[self.layer_idx].keys
+ value_states = curr_past_key_values.layers[self.layer_idx].values
+ else:
+ key_states = self.k_proj(current_states)
+ value_states = self.v_proj(current_states)
+ key_states = key_states.view(*kv_input_shape).transpose(1, 2)
+ value_states = value_states.view(*kv_input_shape).transpose(1, 2)
+
+ if past_key_values is not None:
+ key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx)
+ # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+ if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):
+ past_key_values.is_updated[self.layer_idx] = True
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+class BioGptDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: BioGptConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.embed_dim = config.hidden_size
+
+ self.self_attn = BioGptAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.num_attention_heads,
+ dropout=config.attention_probs_dropout_prob,
+ is_decoder=True,
+ is_causal=True,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.dropout = config.hidden_dropout_prob
+ self.activation_fn = ACT2FN[config.hidden_act]
+ self.activation_dropout = config.activation_dropout
+
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ self.fc1 = nn.Linear(self.embed_dim, config.intermediate_size)
+ self.fc2 = nn.Linear(config.intermediate_size, self.embed_dim)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = True,
+ position_ids: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ past_key_values (`Cache`): cached past key and value projection states
+ """
+ residual = hidden_states
+
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+@auto_docstring
+class BioGptPreTrainedModel(PreTrainedModel):
+ config: BioGptConfig
+ base_model_prefix = "biogpt"
+ supports_gradient_checkpointing = True
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+ _can_compile_fullgraph = True
+ _can_record_outputs = {
+ "hidden_states": BioGptDecoderLayer,
+ "attentions": BioGptAttention,
+ }
+
+
+@auto_docstring
+class BioGptModel(BioGptPreTrainedModel):
+ def __init__(self, config: BioGptConfig):
+ super().__init__(config)
+ self.config = config
+ self.layerdrop = config.layerdrop
+ self.dropout = config.hidden_dropout_prob
+ self.embed_dim = config.hidden_size
+ self.padding_idx = config.pad_token_id
+ embed_scale = math.sqrt(config.hidden_size) if config.scale_embedding else 1.0
+
+ self.embed_tokens = BioGptScaledWordEmbedding(
+ config.vocab_size, self.embed_dim, self.padding_idx, embed_scale=embed_scale
+ )
+ self.embed_positions = BioGptLearnedPositionalEmbedding(config.max_position_embeddings, self.embed_dim)
+
+ self.layers = nn.ModuleList([BioGptDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])
+ self.layer_norm = nn.LayerNorm(self.embed_dim)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ position_ids: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPastAndCrossAttentions:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ # initialize past_key_values
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ batch_size, seq_length = inputs_embeds.size()[:-1]
+ past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+
+ if attention_mask is None:
+ # required mask seq length can be calculated via length of past cache
+ mask_seq_length = past_key_values_length + seq_length
+ attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)
+
+ self_attn_cache = past_key_values
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ input_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=self_attn_cache,
+ )
+
+ # embed positions
+ if position_ids is None:
+ position_ids = torch.arange(seq_length, device=inputs_embeds.device) + past_key_values_length
+ position_ids = position_ids.unsqueeze(0)
+
+ positions = self.embed_positions(attention_mask, past_key_values_length, position_ids=position_ids)
+ hidden_states = inputs_embeds + positions
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ for idx, decoder_layer in enumerate(self.layers):
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop:
+ continue
+
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+ hidden_states = self.layer_norm(hidden_states)
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ BioGPT Model with a `language modeling` head on top for CLM fine-tuning.
+ """
+)
+class BioGptForCausalLM(BioGptPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"output_projection.weight": "biogpt.embed_tokens.weight"}
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.biogpt = BioGptModel(config)
+ self.output_projection = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.output_projection
+
+ def set_output_embeddings(self, new_embeddings):
+ self.output_projection = new_embeddings
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ position_ids: torch.LongTensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithCrossAttentions:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+ `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+ are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+ """
+ outputs = self.biogpt(
+ input_ids,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.output_projection(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithCrossAttentions(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+@auto_docstring
+class BioGptForTokenClassification(BioGptPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.biogpt = BioGptModel(config)
+ if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:
+ classifier_dropout = config.classifier_dropout
+ else:
+ classifier_dropout = config.hidden_dropout_prob
+ self.dropout = nn.Dropout(classifier_dropout)
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ position_ids: torch.LongTensor | None = None,
+ **kwargs,
+ ) -> tuple | TokenClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ transformer_outputs = self.biogpt(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+ hidden_states = transformer_outputs[0]
+ hidden_states = self.dropout(hidden_states)
+ logits = self.classifier(hidden_states)
+
+ loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ if attention_mask is not None:
+ active_loss = attention_mask.view(-1) == 1
+ active_logits = logits.view(-1, self.num_labels)
+ active_labels = torch.where(
+ active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+ )
+ loss = loss_fct(active_logits, active_labels)
+ else:
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+ return TokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The BioGpt Model transformer with a sequence classification head on top (linear layer).
+
+ [`BioGptForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+ (e.g. GPT-2) do.
+
+ Since it does classification on the last token, it is required to know the position of the last token. If a
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+ each row of the batch).
+ """
+)
+class BioGptForSequenceClassification(BioGptPreTrainedModel):
+ def __init__(self, config: BioGptConfig):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.biogpt = BioGptModel(config)
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ position_ids: torch.LongTensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs,
+ ) -> tuple | SequenceClassifierOutputWithPast:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ transformer_outputs = self.biogpt(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ position_ids=position_ids,
+ **kwargs,
+ )
+ hidden_states = transformer_outputs[0]
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.score(hidden_states[:, slice_indices, :])
+
+ if input_ids is not None:
+ batch_size, sequence_length = input_ids.shape[:2]
+ else:
+ batch_size, sequence_length = inputs_embeds.shape[:2]
+
+ if self.config.pad_token_id is None:
+ sequence_length = -1
+ else:
+ if input_ids is not None:
+ sequence_length = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device)
+ else:
+ sequence_length = -1
+ logger.warning_once(
+ f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+ "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+ )
+
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_length]
+
+ loss = None
+ if labels is not None:
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(pooled_logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(pooled_logits, labels)
+
+ return SequenceClassifierOutputWithPast(
+ loss=loss,
+ logits=pooled_logits,
+ past_key_values=transformer_outputs.past_key_values,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+ def get_input_embeddings(self):
+ return self.biogpt.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.biogpt.embed_tokens = value
+
+
+__all__ = [
+ "BioGptForCausalLM",
+ "BioGptForTokenClassification",
+ "BioGptForSequenceClassification",
+ "BioGptModel",
+ "BioGptPreTrainedModel",
+]
diff --git a/third_party/transformers/src/transformers/models/biogpt/modular_biogpt.py b/third_party/transformers/src/transformers/models/biogpt/modular_biogpt.py
new file mode 100644
index 0000000000000000000000000000000000000000..4be4ea1258c467284c3aa86ce6f21f7f235823b7
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/biogpt/modular_biogpt.py
@@ -0,0 +1,509 @@
+# Copyright 2022 The HuggingFace Team and Microsoft Research AI4Science All rights reserved.
+#
+# 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 BioGPT model."""
+
+import math
+
+import torch
+import torch.nn as nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_causal_mask
+from ...modeling_outputs import (
+ BaseModelOutputWithPastAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions,
+ SequenceClassifierOutputWithPast,
+ TokenClassifierOutput,
+)
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import (
+ TransformersKwargs,
+ auto_docstring,
+ can_return_tuple,
+ logger,
+)
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from ..bart.modeling_bart import (
+ BartAttention,
+ BartDecoderLayer,
+ BartScaledWordEmbedding,
+)
+from ..opt.modeling_opt import OPTLearnedPositionalEmbedding
+from .configuration_biogpt import BioGptConfig
+
+
+class BioGptLearnedPositionalEmbedding(OPTLearnedPositionalEmbedding):
+ def forward(
+ self,
+ attention_mask: torch.LongTensor,
+ past_key_values_length: int = 0,
+ position_ids: torch.LongTensor | None = None,
+ ):
+ """`input_ids_shape` is expected to be [bsz x seqlen]."""
+ return super().forward(attention_mask, past_key_values_length, position_ids)
+
+
+class BioGptScaledWordEmbedding(BartScaledWordEmbedding):
+ pass
+
+
+class BioGptAttention(BartAttention):
+ pass
+
+
+class BioGptDecoderLayer(BartDecoderLayer):
+ def __init__(self, config: BioGptConfig, layer_idx: int | None = None):
+ super().__init__(config)
+ self.embed_dim = config.hidden_size
+
+ self.self_attn = BioGptAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.num_attention_heads,
+ dropout=config.attention_probs_dropout_prob,
+ is_decoder=True,
+ is_causal=True,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.dropout = config.hidden_dropout_prob
+ self.activation_fn = ACT2FN[config.hidden_act]
+
+ self.fc1 = nn.Linear(self.embed_dim, config.intermediate_size)
+ self.fc2 = nn.Linear(config.intermediate_size, self.embed_dim)
+
+ del self.encoder_attn
+ del self.encoder_attn_layer_norm
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = True,
+ position_ids: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ past_key_values (`Cache`): cached past key and value projection states
+ """
+ residual = hidden_states
+
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+@auto_docstring
+class BioGptPreTrainedModel(PreTrainedModel):
+ config: BioGptConfig
+ base_model_prefix = "biogpt"
+ supports_gradient_checkpointing = True
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+ _can_compile_fullgraph = True
+ _can_record_outputs = {
+ "hidden_states": BioGptDecoderLayer,
+ "attentions": BioGptAttention,
+ }
+
+
+@auto_docstring
+class BioGptModel(BioGptPreTrainedModel):
+ def __init__(self, config: BioGptConfig):
+ super().__init__(config)
+ self.config = config
+ self.layerdrop = config.layerdrop
+ self.dropout = config.hidden_dropout_prob
+ self.embed_dim = config.hidden_size
+ self.padding_idx = config.pad_token_id
+ embed_scale = math.sqrt(config.hidden_size) if config.scale_embedding else 1.0
+
+ self.embed_tokens = BioGptScaledWordEmbedding(
+ config.vocab_size, self.embed_dim, self.padding_idx, embed_scale=embed_scale
+ )
+ self.embed_positions = BioGptLearnedPositionalEmbedding(config.max_position_embeddings, self.embed_dim)
+
+ self.layers = nn.ModuleList([BioGptDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])
+ self.layer_norm = nn.LayerNorm(self.embed_dim)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ position_ids: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPastAndCrossAttentions:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ # initialize past_key_values
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ batch_size, seq_length = inputs_embeds.size()[:-1]
+ past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+
+ if attention_mask is None:
+ # required mask seq length can be calculated via length of past cache
+ mask_seq_length = past_key_values_length + seq_length
+ attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)
+
+ self_attn_cache = past_key_values
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ input_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=self_attn_cache,
+ )
+
+ # embed positions
+ if position_ids is None:
+ position_ids = torch.arange(seq_length, device=inputs_embeds.device) + past_key_values_length
+ position_ids = position_ids.unsqueeze(0)
+
+ positions = self.embed_positions(attention_mask, past_key_values_length, position_ids=position_ids)
+ hidden_states = inputs_embeds + positions
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ for idx, decoder_layer in enumerate(self.layers):
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop:
+ continue
+
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+ hidden_states = self.layer_norm(hidden_states)
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ BioGPT Model with a `language modeling` head on top for CLM fine-tuning.
+ """
+)
+class BioGptForCausalLM(BioGptPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"output_projection.weight": "biogpt.embed_tokens.weight"}
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.biogpt = BioGptModel(config)
+ self.output_projection = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.output_projection
+
+ def set_output_embeddings(self, new_embeddings):
+ self.output_projection = new_embeddings
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ position_ids: torch.LongTensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithCrossAttentions:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+ `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+ are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+ """
+ outputs = self.biogpt(
+ input_ids,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.output_projection(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithCrossAttentions(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+@auto_docstring
+class BioGptForTokenClassification(BioGptPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.biogpt = BioGptModel(config)
+ if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:
+ classifier_dropout = config.classifier_dropout
+ else:
+ classifier_dropout = config.hidden_dropout_prob
+ self.dropout = nn.Dropout(classifier_dropout)
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ position_ids: torch.LongTensor | None = None,
+ **kwargs,
+ ) -> tuple | TokenClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ transformer_outputs = self.biogpt(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+ hidden_states = transformer_outputs[0]
+ hidden_states = self.dropout(hidden_states)
+ logits = self.classifier(hidden_states)
+
+ loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ if attention_mask is not None:
+ active_loss = attention_mask.view(-1) == 1
+ active_logits = logits.view(-1, self.num_labels)
+ active_labels = torch.where(
+ active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+ )
+ loss = loss_fct(active_logits, active_labels)
+ else:
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+ return TokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The BioGpt Model transformer with a sequence classification head on top (linear layer).
+
+ [`BioGptForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+ (e.g. GPT-2) do.
+
+ Since it does classification on the last token, it is required to know the position of the last token. If a
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+ each row of the batch).
+ """
+)
+class BioGptForSequenceClassification(BioGptPreTrainedModel):
+ def __init__(self, config: BioGptConfig):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.biogpt = BioGptModel(config)
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ position_ids: torch.LongTensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs,
+ ) -> tuple | SequenceClassifierOutputWithPast:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ transformer_outputs = self.biogpt(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ position_ids=position_ids,
+ **kwargs,
+ )
+ hidden_states = transformer_outputs[0]
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.score(hidden_states[:, slice_indices, :])
+
+ if input_ids is not None:
+ batch_size, sequence_length = input_ids.shape[:2]
+ else:
+ batch_size, sequence_length = inputs_embeds.shape[:2]
+
+ if self.config.pad_token_id is None:
+ sequence_length = -1
+ else:
+ if input_ids is not None:
+ sequence_length = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device)
+ else:
+ sequence_length = -1
+ logger.warning_once(
+ f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+ "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+ )
+
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_length]
+
+ loss = None
+ if labels is not None:
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(pooled_logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(pooled_logits, labels)
+
+ return SequenceClassifierOutputWithPast(
+ loss=loss,
+ logits=pooled_logits,
+ past_key_values=transformer_outputs.past_key_values,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+ def get_input_embeddings(self):
+ return self.biogpt.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.biogpt.embed_tokens = value
+
+
+__all__ = [
+ "BioGptForCausalLM",
+ "BioGptForTokenClassification",
+ "BioGptForSequenceClassification",
+ "BioGptModel",
+ "BioGptPreTrainedModel",
+]
diff --git a/third_party/transformers/src/transformers/models/biogpt/tokenization_biogpt.py b/third_party/transformers/src/transformers/models/biogpt/tokenization_biogpt.py
new file mode 100644
index 0000000000000000000000000000000000000000..1bf3e103a0eb837238bdd0f620c7ef423f73387c
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/biogpt/tokenization_biogpt.py
@@ -0,0 +1,329 @@
+# Copyright 2022 The HuggingFace Team and Microsoft Research AI4Science. All rights reserved.
+#
+# 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.
+"""Tokenization classes for BioGPT."""
+
+import json
+import os
+
+from ...tokenization_python import PreTrainedTokenizer
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {
+ "vocab_file": "vocab.json",
+ "merges_file": "merges.txt",
+}
+
+
+def get_pairs(word):
+ """
+ Return set of symbol pairs in a word. word is represented as tuple of symbols (symbols being variable-length
+ strings)
+ """
+ pairs = set()
+ prev_char = word[0]
+ for char in word[1:]:
+ pairs.add((prev_char, char))
+ prev_char = char
+ return pairs
+
+
+class BioGptTokenizer(PreTrainedTokenizer):
+ """
+ Construct an FAIRSEQ Transformer tokenizer. Moses tokenization followed by Byte-Pair Encoding.
+
+ This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+ this superclass for more information regarding those methods.
+
+ Args:
+ vocab_file (`str`):
+ Path to the vocabulary file.
+ merges_file (`str`):
+ Merges file.
+ unk_token (`str`, *optional*, defaults to `""`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ bos_token (`str`, *optional*, defaults to `""`):
+ The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
+
+
+
+ When building a sequence using special tokens, this is not the token that is used for the beginning of
+ sequence. The token used is the `cls_token`.
+
+
+
+ eos_token (`str`, *optional*, defaults to `""`):
+ The end of sequence token.
+
+
+
+ When building a sequence using special tokens, this is not the token that is used for the end of sequence.
+ The token used is the `sep_token`.
+
+
+
+ sep_token (`str`, *optional*, defaults to `""`):
+ The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+ sequence classification or for a text and a question for question answering. It is also used as the last
+ token of a sequence built with special tokens.
+ pad_token (`str`, *optional*, defaults to `""`):
+ The token used for padding, for example when batching sequences of different lengths.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ model_input_names = ["input_ids", "attention_mask"]
+
+ def __init__(
+ self,
+ vocab_file,
+ merges_file,
+ unk_token="",
+ bos_token="",
+ eos_token="",
+ sep_token="",
+ pad_token="",
+ **kwargs,
+ ):
+ try:
+ import sacremoses
+ except ImportError:
+ raise ImportError(
+ "You need to install sacremoses to use BioGptTokenizer. "
+ "See https://pypi.org/project/sacremoses/ for installation."
+ )
+
+ self.lang = "en"
+ self.sm = sacremoses
+ # cache of sm.MosesTokenizer instance
+ self.cache_moses_tokenizer = {}
+ self.cache_moses_detokenizer = {}
+
+ """ Initialisation"""
+ with open(vocab_file, encoding="utf-8") as vocab_handle:
+ self.encoder = json.load(vocab_handle)
+ self.decoder = {v: k for k, v in self.encoder.items()}
+ with open(merges_file, encoding="utf-8") as merges_handle:
+ merges = merges_handle.read().split("\n")[:-1]
+ merges = [tuple(merge.split()[:2]) for merge in merges]
+ self.bpe_ranks = dict(zip(merges, range(len(merges))))
+ self.cache = {}
+
+ super().__init__(
+ bos_token=bos_token,
+ eos_token=eos_token,
+ sep_token=sep_token,
+ unk_token=unk_token,
+ pad_token=pad_token,
+ **kwargs,
+ )
+
+ @property
+ def vocab_size(self):
+ """Returns vocab size"""
+ return len(self.encoder)
+
+ def get_vocab(self):
+ return dict(self.encoder, **self.added_tokens_encoder)
+
+ def moses_tokenize(self, text, lang):
+ if lang not in self.cache_moses_tokenizer:
+ moses_tokenizer = self.sm.MosesTokenizer(lang=lang)
+ self.cache_moses_tokenizer[lang] = moses_tokenizer
+ return self.cache_moses_tokenizer[lang].tokenize(
+ text, aggressive_dash_splits=True, return_str=False, escape=True
+ )
+
+ def moses_detokenize(self, tokens, lang):
+ if lang not in self.cache_moses_detokenizer:
+ moses_detokenizer = self.sm.MosesDetokenizer(lang=lang)
+ self.cache_moses_detokenizer[lang] = moses_detokenizer
+ return self.cache_moses_detokenizer[lang].detokenize(tokens)
+
+ def bpe(self, token):
+ word = tuple(token[:-1]) + (token[-1] + "",)
+ if token in self.cache:
+ return self.cache[token]
+ pairs = get_pairs(word)
+
+ if not pairs:
+ return token + ""
+
+ while True:
+ bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
+ if bigram not in self.bpe_ranks:
+ break
+ first, second = bigram
+ new_word = []
+ i = 0
+ while i < len(word):
+ try:
+ j = word.index(first, i)
+ except ValueError:
+ new_word.extend(word[i:])
+ break
+ else:
+ new_word.extend(word[i:j])
+ i = j
+
+ if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
+ new_word.append(first + second)
+ i += 2
+ else:
+ new_word.append(word[i])
+ i += 1
+ new_word = tuple(new_word)
+ word = new_word
+ if len(word) == 1:
+ break
+ else:
+ pairs = get_pairs(word)
+ word = " ".join(word)
+ if word == "\n ":
+ word = "\n"
+ self.cache[token] = word
+ return word
+
+ def _tokenize(self, text, bypass_tokenizer=False):
+ """Returns a tokenized string."""
+ if bypass_tokenizer:
+ text = text.split()
+ else:
+ text = self.moses_tokenize(text, self.lang)
+
+ split_tokens = []
+ for token in text:
+ if token:
+ split_tokens.extend(list(self.bpe(token).split(" ")))
+
+ return split_tokens
+
+ def _convert_token_to_id(self, token):
+ """Converts a token (str) in an id using the vocab."""
+ return self.encoder.get(token, self.encoder.get(self.unk_token))
+
+ def _convert_id_to_token(self, index):
+ """Converts an index (integer) in a token (str) using the vocab."""
+ return self.decoder.get(index, self.unk_token)
+
+ def convert_tokens_to_string(self, tokens):
+ """Converts a sequence of tokens (string) in a single string."""
+ # remove BPE
+ tokens = [t.replace(" ", "").replace("", " ") for t in tokens]
+ tokens = "".join(tokens).split()
+ # detokenize
+ text = self.moses_detokenize(tokens, self.lang)
+ return text
+
+ def build_inputs_with_special_tokens(
+ self, token_ids_0: list[int], token_ids_1: list[int] | None = None
+ ) -> list[int]:
+ """
+ Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
+ adding special tokens. A BioGPT sequence has the following format:
+
+ - single sequence: ` X `
+ - pair of sequences: ` A B `
+
+ Args:
+ token_ids_0 (`List[int]`):
+ List of IDs to which the special tokens will be added.
+ token_ids_1 (`List[int]`, *optional*):
+ Optional second list of IDs for sequence pairs.
+
+ Returns:
+ `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
+ """
+ if token_ids_1 is None:
+ return [self.sep_token_id] + token_ids_0
+ sep = [self.sep_token_id]
+ return sep + token_ids_0 + sep + token_ids_1
+
+ def get_special_tokens_mask(
+ self, token_ids_0: list[int], token_ids_1: list[int] | None = None, already_has_special_tokens: bool = False
+ ) -> list[int]:
+ """
+ Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
+ special tokens using the tokenizer `prepare_for_model` method.
+
+ Args:
+ token_ids_0 (`List[int]`):
+ List of IDs.
+ token_ids_1 (`List[int]`, *optional*):
+ Optional second list of IDs for sequence pairs.
+ already_has_special_tokens (`bool`, *optional*, defaults to `False`):
+ Whether or not the token list is already formatted with special tokens for the model.
+
+ Returns:
+ `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
+ """
+ if already_has_special_tokens:
+ return super().get_special_tokens_mask(
+ token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
+ )
+ # no bos used in fairseq
+ if token_ids_1 is not None:
+ return [1] + ([0] * len(token_ids_0)) + [1] + ([0] * len(token_ids_1))
+ return [1] + ([0] * len(token_ids_0))
+
+ def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
+ if not os.path.isdir(save_directory):
+ logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+ return
+ vocab_file = os.path.join(
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+ )
+ merge_file = os.path.join(
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
+ )
+
+ with open(vocab_file, "w", encoding="utf-8") as f:
+ f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
+
+ index = 0
+ with open(merge_file, "w", encoding="utf-8") as writer:
+ for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
+ if index != token_index:
+ logger.warning(
+ f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
+ " Please check that the tokenizer is not corrupted!"
+ )
+ index = token_index
+ writer.write(" ".join(bpe_tokens) + "\n")
+ index += 1
+
+ return vocab_file, merge_file
+
+ def __getstate__(self):
+ state = self.__dict__.copy()
+ state["sm"] = None
+ return state
+
+ def __setstate__(self, d):
+ self.__dict__ = d
+
+ try:
+ import sacremoses
+ except ImportError:
+ raise ImportError(
+ "You need to install sacremoses to use XLMTokenizer. "
+ "See https://pypi.org/project/sacremoses/ for installation."
+ )
+
+ self.sm = sacremoses
+
+
+__all__ = ["BioGptTokenizer"]
diff --git a/third_party/transformers/src/transformers/models/bit/__init__.py b/third_party/transformers/src/transformers/models/bit/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a482bb790ea9fc7f5a2ecc21cd0951ebd572e314
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bit/__init__.py
@@ -0,0 +1,29 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_bit import *
+ from .image_processing_bit import *
+ from .image_processing_pil_bit import *
+ from .modeling_bit import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/bit/configuration_bit.py b/third_party/transformers/src/transformers/models/bit/configuration_bit.py
new file mode 100644
index 0000000000000000000000000000000000000000..8b81f51ee739a52bc7714ea892ccabe438fb83b2
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bit/configuration_bit.py
@@ -0,0 +1,95 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""BiT model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...backbone_utils import BackboneConfigMixin
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="google/bit-50")
+@strict
+class BitConfig(BackboneConfigMixin, PreTrainedConfig):
+ r"""
+ layer_type (`str`, *optional*, defaults to `"preactivation"`):
+ The layer to use, it can be either `"preactivation"` or `"bottleneck"`.
+ global_padding (`str`, *optional*):
+ Padding strategy to use for the convolutional layers. Can be either `"valid"`, `"same"`, or `None`.
+ num_groups (`int`, *optional*, defaults to 32):
+ Number of groups used for the `BitGroupNormActivation` layers.
+ embedding_dynamic_padding (`bool`, *optional*, defaults to `False`):
+ Whether or not to make use of dynamic padding for the embedding layer.
+ width_factor (`int`, *optional*, defaults to 1):
+ The width factor for the model.
+
+ Example:
+ ```python
+ >>> from transformers import BitConfig, BitModel
+
+ >>> # Initializing a BiT bit-50 style configuration
+ >>> configuration = BitConfig()
+
+ >>> # Initializing a model (with random weights) from the bit-50 style configuration
+ >>> model = BitModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+ """
+
+ model_type = "bit"
+ layer_types = ["preactivation", "bottleneck"]
+ supported_padding = [None, "SAME", "VALID"]
+
+ num_channels: int = 3
+ embedding_size: int = 64
+ hidden_sizes: list[int] | tuple[int, ...] = (256, 512, 1024, 2048)
+ depths: list[int] | tuple[int, ...] = (3, 4, 6, 3)
+ layer_type: str = "preactivation"
+ hidden_act: str = "relu"
+ global_padding: str | None = None
+ num_groups: int = 32
+ drop_path_rate: float | int = 0.0
+ embedding_dynamic_padding: bool = False
+ output_stride: int = 32
+ width_factor: int = 1
+ _out_features: list[str] | None = None
+ _out_indices: list[int] | None = None
+
+ def __post_init__(self, **kwargs):
+ self.hidden_sizes = list(self.hidden_sizes)
+ self.depths = list(self.depths)
+
+ if self.global_padding is not None:
+ self.global_padding = self.global_padding.upper()
+
+ self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, len(self.depths) + 1)]
+ self.set_output_features_output_indices(
+ out_indices=kwargs.pop("out_indices", None), out_features=kwargs.pop("out_features", None)
+ )
+
+ super().__post_init__(**kwargs)
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ if self.layer_type not in self.layer_types:
+ raise ValueError(f"layer_type={self.layer_type} is not one of {','.join(self.layer_types)}")
+
+ if self.global_padding not in self.supported_padding:
+ raise ValueError(f"Padding strategy {self.global_padding} not supported")
+
+
+__all__ = ["BitConfig"]
diff --git a/third_party/transformers/src/transformers/models/bit/convert_bit_to_pytorch.py b/third_party/transformers/src/transformers/models/bit/convert_bit_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..d2b300820d160c2f6d275fd82fcf2b1b0e7fb98e
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bit/convert_bit_to_pytorch.py
@@ -0,0 +1,178 @@
+# Copyright 2022 The HuggingFace Inc. team.
+#
+# 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.
+"""Convert BiT checkpoints from the timm library."""
+
+import argparse
+import json
+from io import BytesIO
+from pathlib import Path
+
+import httpx
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image
+from timm import create_model
+from timm.data import resolve_data_config
+from timm.data.transforms_factory import create_transform
+
+from transformers import BitConfig, BitForImageClassification, BitImageProcessor
+from transformers.image_utils import PILImageResampling
+from transformers.utils import logging
+
+
+logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+
+def get_config(model_name):
+ repo_id = "huggingface/label-files"
+ filename = "imagenet-1k-id2label.json"
+ id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
+ id2label = {int(k): v for k, v in id2label.items()}
+ label2id = {v: k for k, v in id2label.items()}
+
+ conv_layer = "std_conv" if "bit" in model_name else False
+
+ # note that when using BiT as backbone for ViT-hybrid checkpoints,
+ # one needs to additionally set config.layer_type = "bottleneck", config.stem_type = "same",
+ # config.conv_layer = "std_conv_same"
+ config = BitConfig(
+ conv_layer=conv_layer,
+ num_labels=1000,
+ id2label=id2label,
+ label2id=label2id,
+ )
+
+ return config
+
+
+def rename_key(name):
+ if "stem.conv" in name:
+ name = name.replace("stem.conv", "bit.embedder.convolution")
+ if "blocks" in name:
+ name = name.replace("blocks", "layers")
+ if "head.fc" in name:
+ name = name.replace("head.fc", "classifier.1")
+ if name.startswith("norm"):
+ name = "bit." + name
+ if "bit" not in name and "classifier" not in name:
+ name = "bit.encoder." + name
+
+ return name
+
+
+# We will verify our results on an image of cute cats
+def prepare_img():
+ url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ with httpx.stream("GET", url) as response:
+ image = Image.open(BytesIO(response.read()))
+ return image
+
+
+@torch.no_grad()
+def convert_bit_checkpoint(model_name, pytorch_dump_folder_path, push_to_hub=False):
+ """
+ Copy/paste/tweak model's weights to our BiT structure.
+ """
+
+ # define default BiT configuration
+ config = get_config(model_name)
+
+ # load original model from timm
+ timm_model = create_model(model_name, pretrained=True)
+ timm_model.eval()
+
+ # load state_dict of original model
+ state_dict = timm_model.state_dict()
+ for key in state_dict.copy():
+ val = state_dict.pop(key)
+ state_dict[rename_key(key)] = val.squeeze() if "head" in key else val
+
+ # load HuggingFace model
+ model = BitForImageClassification(config)
+ model.eval()
+ model.load_state_dict(state_dict)
+
+ # create image processor
+ transform = create_transform(**resolve_data_config({}, model=timm_model))
+ timm_transforms = transform.transforms
+
+ pillow_resamplings = {
+ "bilinear": PILImageResampling.BILINEAR,
+ "bicubic": PILImageResampling.BICUBIC,
+ "nearest": PILImageResampling.NEAREST,
+ }
+
+ processor = BitImageProcessor(
+ do_resize=True,
+ size={"shortest_edge": timm_transforms[0].size},
+ resample=pillow_resamplings[timm_transforms[0].interpolation.value],
+ do_center_crop=True,
+ crop_size={"height": timm_transforms[1].size[0], "width": timm_transforms[1].size[1]},
+ do_normalize=True,
+ image_mean=timm_transforms[-1].mean.tolist(),
+ image_std=timm_transforms[-1].std.tolist(),
+ )
+
+ image = prepare_img()
+ timm_pixel_values = transform(image).unsqueeze(0)
+ pixel_values = processor(image, return_tensors="pt").pixel_values
+
+ # verify pixel values
+ assert torch.allclose(timm_pixel_values, pixel_values)
+
+ # verify logits
+ with torch.no_grad():
+ outputs = model(pixel_values)
+ logits = outputs.logits
+
+ print("Logits:", logits[0, :3])
+ print("Predicted class:", model.config.id2label[logits.argmax(-1).item()])
+ timm_logits = timm_model(pixel_values)
+ assert timm_logits.shape == outputs.logits.shape
+ assert torch.allclose(timm_logits, outputs.logits, atol=1e-3)
+ print("Looks ok!")
+
+ if pytorch_dump_folder_path is not None:
+ Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
+ print(f"Saving model {model_name} and processor to {pytorch_dump_folder_path}")
+ model.save_pretrained(pytorch_dump_folder_path)
+ processor.save_pretrained(pytorch_dump_folder_path)
+
+ if push_to_hub:
+ print(f"Pushing model {model_name} and processor to the hub")
+ model.push_to_hub(f"ybelkada/{model_name}")
+ processor.push_to_hub(f"ybelkada/{model_name}")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ # Required parameters
+ parser.add_argument(
+ "--model_name",
+ default="resnetv2_50x1_bitm",
+ type=str,
+ help="Name of the BiT timm model you'd like to convert.",
+ )
+ parser.add_argument(
+ "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
+ )
+ parser.add_argument(
+ "--push_to_hub",
+ action="store_true",
+ help="Whether to push the model to the hub.",
+ )
+
+ args = parser.parse_args()
+ convert_bit_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
diff --git a/third_party/transformers/src/transformers/models/bit/image_processing_bit.py b/third_party/transformers/src/transformers/models/bit/image_processing_bit.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c2eb87272fda9071878195428a64dc9da739832
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bit/image_processing_bit.py
@@ -0,0 +1,36 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Image processor class for BiT."""
+
+from ...image_processing_backends import TorchvisionBackend
+from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, PILImageResampling
+from ...utils import auto_docstring
+
+
+@auto_docstring
+class BitImageProcessor(TorchvisionBackend):
+ resample = PILImageResampling.BICUBIC
+ image_mean = OPENAI_CLIP_MEAN
+ image_std = OPENAI_CLIP_STD
+ size = {"shortest_edge": 224}
+ default_to_square = False
+ crop_size = {"height": 224, "width": 224}
+ do_resize = True
+ do_center_crop = True
+ do_rescale = True
+ do_normalize = True
+ do_convert_rgb = True
+
+
+__all__ = ["BitImageProcessor"]
diff --git a/third_party/transformers/src/transformers/models/bit/image_processing_pil_bit.py b/third_party/transformers/src/transformers/models/bit/image_processing_pil_bit.py
new file mode 100644
index 0000000000000000000000000000000000000000..5e355175cefd9df2765c145b410bae36cde1c9bf
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bit/image_processing_pil_bit.py
@@ -0,0 +1,36 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Image processor class for BiT."""
+
+from ...image_processing_backends import PilBackend
+from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, PILImageResampling
+from ...utils import auto_docstring
+
+
+@auto_docstring
+class BitImageProcessorPil(PilBackend):
+ resample = PILImageResampling.BICUBIC
+ image_mean = OPENAI_CLIP_MEAN
+ image_std = OPENAI_CLIP_STD
+ size = {"shortest_edge": 224}
+ default_to_square = False
+ crop_size = {"height": 224, "width": 224}
+ do_resize = True
+ do_center_crop = True
+ do_rescale = True
+ do_normalize = True
+ do_convert_rgb = True
+
+
+__all__ = ["BitImageProcessorPil"]
diff --git a/third_party/transformers/src/transformers/models/bit/modeling_bit.py b/third_party/transformers/src/transformers/models/bit/modeling_bit.py
new file mode 100644
index 0000000000000000000000000000000000000000..e2056f057f4a7d6a238aee148e57fa740f6d75a7
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bit/modeling_bit.py
@@ -0,0 +1,832 @@
+# Copyright 2022 Google AI and The HuggingFace Inc. team. All rights reserved.
+#
+# 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 BiT model. Also supports backbone for ViT hybrid."""
+
+import collections
+import math
+
+import numpy as np
+import torch
+from torch import Tensor, nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...backbone_utils import BackboneMixin, filter_output_hidden_states
+from ...modeling_outputs import (
+ BackboneOutput,
+ BaseModelOutputWithNoAttention,
+ BaseModelOutputWithPoolingAndNoAttention,
+ ImageClassifierOutputWithNoAttention,
+)
+from ...modeling_utils import PreTrainedModel
+from ...utils import auto_docstring, logging
+from ...utils.generic import can_return_tuple
+from .configuration_bit import BitConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+def get_padding_value(padding=None, kernel_size=7, stride=1, dilation=1) -> tuple[tuple, bool]:
+ r"""
+ Utility function to get the tuple padding value given the kernel_size and padding.
+
+ Args:
+ padding (Union[`str`, `int`], *optional*):
+ Padding value, can be either `"same"`, `"valid"`. If a different value is provided the default padding from
+ PyTorch is used.
+ kernel_size (`int`, *optional*, defaults to 7):
+ Kernel size of the convolution layers.
+ stride (`int`, *optional*, defaults to 1):
+ Stride value of the convolution layers.
+ dilation (`int`, *optional*, defaults to 1):
+ Dilation value of the convolution layers.
+ """
+ dynamic = False
+ if padding is None:
+ padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
+ return padding, dynamic
+
+ if isinstance(padding, str):
+ # for any string padding, the padding will be calculated for you, one of three ways
+ padding = padding.lower()
+ if padding == "same":
+ # TF compatible 'SAME' padding, has a performance and GPU memory allocation impact
+ if stride == 1 and (dilation * (kernel_size - 1)) % 2 == 0:
+ # static case, no extra overhead
+ padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
+ else:
+ # dynamic 'SAME' padding, has runtime/GPU memory overhead
+ padding = 0
+ dynamic = True
+ elif padding == "valid":
+ # 'VALID' padding, same as padding=0
+ padding = 0
+ else:
+ # Default to PyTorch style 'same'-ish symmetric padding
+ padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
+ return padding, dynamic
+
+
+class WeightStandardizedConv2d(nn.Conv2d):
+ """Conv2d with Weight Standardization. Used for ViT Hybrid model.
+
+ Paper: [Micro-Batch Training with Batch-Channel Normalization and Weight
+ Standardization](https://huggingface.co/papers/1903.10520)
+ """
+
+ def __init__(
+ self,
+ in_channel,
+ out_channels,
+ kernel_size,
+ stride=1,
+ padding="SAME",
+ dilation=1,
+ groups=1,
+ bias=False,
+ eps=1e-6,
+ ):
+ padding, is_dynamic = get_padding_value(padding, kernel_size, stride=stride, dilation=dilation)
+ super().__init__(
+ in_channel,
+ out_channels,
+ kernel_size,
+ stride=stride,
+ padding=padding,
+ dilation=dilation,
+ groups=groups,
+ bias=bias,
+ )
+ if is_dynamic:
+ self.pad = DynamicPad2d(kernel_size, stride, dilation)
+ else:
+ self.pad = None
+ self.eps = eps
+
+ def forward(self, hidden_state):
+ if self.pad is not None:
+ hidden_state = self.pad(hidden_state)
+ weight = nn.functional.batch_norm(
+ self.weight.reshape(1, self.out_channels, -1), None, None, training=True, momentum=0.0, eps=self.eps
+ ).reshape_as(self.weight)
+ hidden_state = nn.functional.conv2d(
+ hidden_state, weight, self.bias, self.stride, self.padding, self.dilation, self.groups
+ )
+ return hidden_state
+
+
+class BitGroupNormActivation(nn.GroupNorm):
+ r"""
+ A module that combines group normalization with an activation function.
+ """
+
+ def __init__(self, config, num_channels, eps=1e-5, affine=True, apply_activation=True):
+ super().__init__(config.num_groups, num_channels, eps=eps, affine=affine)
+ if apply_activation:
+ self.activation = ACT2FN[config.hidden_act]
+ else:
+ self.activation = nn.Identity()
+
+ def forward(self, hidden_state):
+ hidden_state = nn.functional.group_norm(hidden_state, self.num_groups, self.weight, self.bias, self.eps)
+ hidden_state = self.activation(hidden_state)
+ return hidden_state
+
+
+class DynamicPad2d(nn.Module):
+ r"""
+ A module that wraps dynamic padding of any input, given the parameters of the convolutional layer and the input
+ hidden states.
+ """
+
+ def __init__(self, kernel_size, stride, dilation, value=0):
+ super().__init__()
+ # Safety checkers
+ if isinstance(kernel_size, int):
+ kernel_size = (kernel_size, kernel_size)
+
+ if isinstance(stride, int):
+ stride = (stride, stride)
+
+ if isinstance(dilation, int):
+ dilation = (dilation, dilation)
+
+ self.kernel_size = kernel_size
+ self.stride = stride
+ self.dilation = dilation
+ self.value = value
+
+ def compute_padding(x, kernel_size, stride, dilation):
+ return max((math.ceil(x / stride) - 1) * stride + (kernel_size - 1) * dilation + 1 - x, 0)
+
+ self.compute_padding = compute_padding
+
+ def forward(self, input):
+ # Get width and height
+ input_height, input_width = input.size()[-2:]
+
+ # Compute the padding values
+ padding_height = self.compute_padding(input_height, self.kernel_size[0], self.stride[0], self.dilation[0])
+ padding_width = self.compute_padding(input_width, self.kernel_size[1], self.stride[1], self.dilation[1])
+
+ # apply pad
+ if padding_height > 0 or padding_width > 0:
+ input = nn.functional.pad(
+ input,
+ [
+ padding_width // 2,
+ padding_width - padding_width // 2,
+ padding_height // 2,
+ padding_height - padding_height // 2,
+ ],
+ value=self.value,
+ )
+ return input
+
+
+class BitMaxPool2d(nn.MaxPool2d):
+ def __init__(
+ self,
+ kernel_size: int,
+ stride=None,
+ dilation=1,
+ ceil_mode=False,
+ padding=(0, 0),
+ padding_value=0,
+ use_dynamic_padding=True,
+ ):
+ kernel_size = kernel_size if isinstance(kernel_size, collections.abc.Iterable) else (kernel_size, kernel_size)
+ stride = stride if isinstance(stride, collections.abc.Iterable) else (stride, stride)
+ dilation = dilation if isinstance(dilation, collections.abc.Iterable) else (dilation, dilation)
+ super().__init__(kernel_size, stride, padding, dilation, ceil_mode)
+ if use_dynamic_padding:
+ self.pad = DynamicPad2d(kernel_size, stride, dilation, padding_value)
+ else:
+ self.pad = nn.Identity()
+
+ def forward(self, hidden_states):
+ hidden_states = self.pad(hidden_states)
+ return nn.functional.max_pool2d(
+ hidden_states, self.kernel_size, self.stride, self.padding, self.dilation, self.ceil_mode
+ )
+
+
+class BitEmbeddings(nn.Module):
+ """
+ BiT Embeddings (stem) composed of a single aggressive convolution.
+ """
+
+ def __init__(self, config: BitConfig):
+ super().__init__()
+
+ self.convolution = WeightStandardizedConv2d(
+ config.num_channels,
+ config.embedding_size,
+ kernel_size=7,
+ stride=2,
+ eps=1e-8,
+ padding=config.global_padding,
+ )
+
+ self.pooler = BitMaxPool2d(kernel_size=3, stride=2, use_dynamic_padding=config.embedding_dynamic_padding)
+
+ # Use the same padding strategy as convolutional layers
+ if config.global_padding is not None and config.global_padding.upper() == "SAME":
+ self.pad = nn.Identity()
+ else:
+ self.pad = nn.ConstantPad2d(padding=(1, 1, 1, 1), value=0.0)
+
+ if config.layer_type != "preactivation":
+ self.norm = BitGroupNormActivation(config, num_channels=config.embedding_size)
+ else:
+ self.norm = nn.Identity()
+
+ self.num_channels = config.num_channels
+
+ def forward(self, pixel_values: Tensor) -> Tensor:
+ num_channels = pixel_values.shape[1]
+ if num_channels != self.num_channels:
+ raise ValueError(
+ "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+ )
+
+ embedding = self.convolution(pixel_values)
+
+ embedding = self.pad(embedding)
+
+ embedding = self.norm(embedding)
+
+ embedding = self.pooler(embedding)
+
+ return embedding
+
+
+# Copied from transformers.models.convnext.modeling_convnext.drop_path
+def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
+ """
+ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+ """
+ if drop_prob == 0.0 or not training:
+ return input
+ keep_prob = 1 - drop_prob
+ shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
+ random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
+ random_tensor.floor_() # binarize
+ output = input.div(keep_prob) * random_tensor
+ return output
+
+
+# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->Bit
+class BitDropPath(nn.Module):
+ """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+ def __init__(self, drop_prob: float | None = None) -> None:
+ super().__init__()
+ self.drop_prob = drop_prob
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ return drop_path(hidden_states, self.drop_prob, self.training)
+
+ def extra_repr(self) -> str:
+ return f"p={self.drop_prob}"
+
+
+def make_div(value, divisor=8):
+ min_value = divisor
+ new_value = max(min_value, int(value + divisor / 2) // divisor * divisor)
+ if new_value < 0.9 * value:
+ new_value += divisor
+ return new_value
+
+
+class BitPreActivationBottleneckLayer(nn.Module):
+ """Pre-activation (v2) bottleneck block.
+ Follows the implementation of "Identity Mappings in Deep Residual Networks":
+ https://github.com/KaimingHe/resnet-1k-layers/blob/master/resnet-pre-act.lua
+
+ Except it puts the stride on 3x3 conv when available.
+ """
+
+ def __init__(
+ self,
+ config,
+ in_channels,
+ out_channels=None,
+ bottle_ratio=0.25,
+ stride=1,
+ dilation=1,
+ first_dilation=None,
+ groups=1,
+ drop_path_rate=0.0,
+ is_first_layer=False,
+ ):
+ super().__init__()
+
+ first_dilation = first_dilation or dilation
+
+ out_channels = out_channels or in_channels
+ mid_channels = make_div(out_channels * bottle_ratio)
+
+ if is_first_layer:
+ self.downsample = BitDownsampleConv(
+ config,
+ in_channels,
+ out_channels,
+ stride=stride,
+ preact=True,
+ )
+ else:
+ self.downsample = None
+
+ self.norm1 = BitGroupNormActivation(config, in_channels)
+ self.conv1 = WeightStandardizedConv2d(in_channels, mid_channels, 1, eps=1e-8, padding=config.global_padding)
+
+ self.norm2 = BitGroupNormActivation(config, num_channels=mid_channels)
+ self.conv2 = WeightStandardizedConv2d(
+ mid_channels, mid_channels, 3, stride=stride, groups=groups, eps=1e-8, padding=config.global_padding
+ )
+
+ self.norm3 = BitGroupNormActivation(config, mid_channels)
+ self.conv3 = WeightStandardizedConv2d(mid_channels, out_channels, 1, eps=1e-8, padding=config.global_padding)
+
+ self.drop_path = BitDropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
+
+ def forward(self, hidden_states):
+ hidden_states_preact = self.norm1(hidden_states)
+
+ # shortcut branch
+ shortcut = hidden_states
+ if self.downsample is not None:
+ shortcut = self.downsample(hidden_states_preact)
+
+ # residual branch
+ hidden_states = self.conv1(hidden_states_preact)
+ hidden_states = self.conv2(self.norm2(hidden_states))
+ hidden_states = self.conv3(self.norm3(hidden_states))
+ hidden_states = self.drop_path(hidden_states)
+ return hidden_states + shortcut
+
+
+class BitBottleneckLayer(nn.Module):
+ """Non Pre-activation bottleneck block, equivalent to V1.5/V1b bottleneck. Used for ViT Hybrid."""
+
+ def __init__(
+ self,
+ config,
+ in_channels,
+ out_channels=None,
+ bottle_ratio=0.25,
+ stride=1,
+ dilation=1,
+ first_dilation=None,
+ groups=1,
+ drop_path_rate=0.0,
+ is_first_layer=False,
+ ):
+ super().__init__()
+ first_dilation = first_dilation or dilation
+
+ out_channels = out_channels or in_channels
+ mid_chs = make_div(out_channels * bottle_ratio)
+
+ if is_first_layer:
+ self.downsample = BitDownsampleConv(
+ config,
+ in_channels,
+ out_channels,
+ stride=stride,
+ preact=False,
+ )
+ else:
+ self.downsample = None
+
+ self.conv1 = WeightStandardizedConv2d(in_channels, mid_chs, 1, eps=1e-8, padding=config.global_padding)
+ self.norm1 = BitGroupNormActivation(config, num_channels=mid_chs)
+ self.conv2 = WeightStandardizedConv2d(
+ mid_chs,
+ mid_chs,
+ 3,
+ stride=stride,
+ dilation=first_dilation,
+ groups=groups,
+ eps=1e-8,
+ padding=config.global_padding,
+ )
+ self.norm2 = BitGroupNormActivation(config, num_channels=mid_chs)
+ self.conv3 = WeightStandardizedConv2d(mid_chs, out_channels, 1, eps=1e-8, padding=config.global_padding)
+ self.norm3 = BitGroupNormActivation(config, num_channels=out_channels, apply_activation=False)
+ self.drop_path = BitDropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
+
+ self.activation = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_states):
+ # shortcut branch
+ shortcut = hidden_states
+ if self.downsample is not None:
+ shortcut = self.downsample(hidden_states)
+
+ # residual
+ hidden_states = self.conv1(hidden_states)
+ hidden_states = self.norm1(hidden_states)
+
+ hidden_states = self.conv2(hidden_states)
+ hidden_states = self.norm2(hidden_states)
+
+ hidden_states = self.conv3(hidden_states)
+ hidden_states = self.norm3(hidden_states)
+
+ hidden_states = self.drop_path(hidden_states)
+ hidden_states = self.activation(hidden_states + shortcut)
+ return hidden_states
+
+
+class BitDownsampleConv(nn.Module):
+ def __init__(
+ self,
+ config,
+ in_channels,
+ out_channels,
+ stride=1,
+ preact=True,
+ ):
+ super().__init__()
+ self.conv = WeightStandardizedConv2d(
+ in_channels, out_channels, 1, stride=stride, eps=1e-8, padding=config.global_padding
+ )
+ self.norm = (
+ nn.Identity()
+ if preact
+ else BitGroupNormActivation(config, num_channels=out_channels, apply_activation=False)
+ )
+
+ def forward(self, x):
+ return self.norm(self.conv(x))
+
+
+class BitStage(nn.Module):
+ """
+ A ResNet v2 stage composed by stacked layers.
+ """
+
+ def __init__(
+ self,
+ config,
+ in_channels,
+ out_channels,
+ stride,
+ dilation,
+ depth,
+ bottle_ratio=0.25,
+ layer_dropout=None,
+ ):
+ super().__init__()
+
+ first_dilation = 1 if dilation in (1, 2) else 2
+
+ # Get the layer type
+ if config.layer_type == "bottleneck":
+ layer_cls = BitBottleneckLayer
+ else:
+ layer_cls = BitPreActivationBottleneckLayer
+
+ prev_chs = in_channels
+ self.layers = nn.Sequential()
+ for layer_idx in range(depth):
+ # Get the current hyper-parameters
+ stride, drop_path_rate, is_first_layer = self._get_updated_hyperparameters(
+ layer_idx, stride, layer_dropout
+ )
+
+ self.layers.add_module(
+ str(layer_idx),
+ layer_cls(
+ config,
+ prev_chs,
+ out_channels,
+ stride=stride,
+ dilation=dilation,
+ bottle_ratio=bottle_ratio,
+ first_dilation=first_dilation,
+ drop_path_rate=drop_path_rate,
+ is_first_layer=is_first_layer,
+ ),
+ )
+ prev_chs = out_channels
+ first_dilation = dilation
+
+ def _get_updated_hyperparameters(self, layer_idx, stride, layer_dropout):
+ r"""
+ Get the new hyper-parameters with respect to the previous ones and the index of the current layer.
+ """
+ if layer_dropout:
+ drop_path_rate = layer_dropout[layer_idx]
+ else:
+ drop_path_rate = 0.0
+
+ if layer_idx != 0:
+ stride = 1
+
+ is_first_layer = layer_idx == 0
+
+ return stride, drop_path_rate, is_first_layer
+
+ def forward(self, input: Tensor) -> Tensor:
+ hidden_state = input
+ for _, layer in enumerate(self.layers):
+ hidden_state = layer(hidden_state)
+ return hidden_state
+
+
+class BitEncoder(nn.Module):
+ def __init__(self, config: BitConfig):
+ super().__init__()
+ self.stages = nn.ModuleList([])
+
+ prev_chs = config.embedding_size
+
+ # These needs to stay hardcoded
+ current_stride = 4
+ dilation = 1
+
+ layer_dropouts = [
+ x.tolist()
+ for x in torch.Tensor(np.linspace(0, config.drop_path_rate, sum(config.depths))).split(config.depths)
+ ]
+
+ for stage_idx, (current_depth, current_hidden_size, layer_dropout) in enumerate(
+ zip(config.depths, config.hidden_sizes, layer_dropouts)
+ ):
+ # Get the updated hyper params
+ out_channels, stride, dilation = self._get_updated_hyperparameters(
+ stage_idx, current_stride, current_hidden_size, dilation, config
+ )
+
+ stage = BitStage(
+ config,
+ prev_chs,
+ out_channels,
+ stride=stride,
+ dilation=dilation,
+ depth=current_depth,
+ layer_dropout=layer_dropout,
+ )
+
+ prev_chs = out_channels
+ current_stride *= stride
+
+ self.stages.add_module(str(stage_idx), stage)
+
+ def _get_updated_hyperparameters(self, stage_idx, current_stride, current_hidden_size, dilation, config):
+ out_channels = make_div(current_hidden_size * config.width_factor)
+ stride = 1 if stage_idx == 0 else 2
+ if current_stride >= config.output_stride:
+ dilation *= stride
+ stride = 1
+ return out_channels, stride, dilation
+
+ def forward(
+ self, hidden_state: Tensor, output_hidden_states: bool = False, return_dict: bool = True
+ ) -> BaseModelOutputWithNoAttention:
+ hidden_states = () if output_hidden_states else None
+
+ for stage_module in self.stages:
+ if output_hidden_states:
+ hidden_states = hidden_states + (hidden_state,)
+
+ hidden_state = stage_module(hidden_state)
+
+ if output_hidden_states:
+ hidden_states = hidden_states + (hidden_state,)
+
+ if not return_dict:
+ return tuple(v for v in [hidden_state, hidden_states] if v is not None)
+
+ return BaseModelOutputWithNoAttention(
+ last_hidden_state=hidden_state,
+ hidden_states=hidden_states,
+ )
+
+
+@auto_docstring
+class BitPreTrainedModel(PreTrainedModel):
+ config: BitConfig
+ base_model_prefix = "bit"
+ input_modalities = ("image",)
+ main_input_name = "pixel_values"
+ _no_split_modules = ["BitEmbeddings"]
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ if isinstance(module, nn.Conv2d):
+ init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
+ # copied from the `reset_parameters` method of `class Linear(Module)` in `torch`.
+ elif isinstance(module, nn.Linear):
+ init.kaiming_uniform_(module.weight, a=math.sqrt(5))
+ if module.bias is not None:
+ fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(module.weight)
+ bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+ init.uniform_(module.bias, -bound, bound)
+ elif isinstance(module, (nn.BatchNorm2d, nn.GroupNorm)):
+ init.constant_(module.weight, 1)
+ init.constant_(module.bias, 0)
+ if getattr(module, "running_mean", None) is not None:
+ init.zeros_(module.running_mean)
+ init.ones_(module.running_var)
+ init.zeros_(module.num_batches_tracked)
+
+
+@auto_docstring
+class BitModel(BitPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.config = config
+
+ self.embedder = BitEmbeddings(config)
+
+ self.encoder = BitEncoder(config)
+ self.norm = (
+ BitGroupNormActivation(config, num_channels=config.hidden_sizes[-1])
+ if config.layer_type == "preactivation"
+ else nn.Identity()
+ )
+
+ self.pooler = nn.AdaptiveAvgPool2d((1, 1))
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: Tensor,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> BaseModelOutputWithPoolingAndNoAttention:
+ 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.return_dict
+
+ embedding_output = self.embedder(pixel_values)
+
+ encoder_outputs = self.encoder(
+ embedding_output, output_hidden_states=output_hidden_states, return_dict=return_dict
+ )
+
+ last_hidden_state = encoder_outputs[0]
+
+ last_hidden_state = self.norm(last_hidden_state)
+
+ pooled_output = self.pooler(last_hidden_state)
+
+ if not return_dict:
+ return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPoolingAndNoAttention(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ BiT Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
+ ImageNet.
+ """
+)
+class BitForImageClassification(BitPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.bit = BitModel(config)
+ # classification head
+ self.classifier = nn.Sequential(
+ nn.Flatten(),
+ nn.Linear(config.hidden_sizes[-1], config.num_labels) if config.num_labels > 0 else nn.Identity(),
+ )
+ # initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> ImageClassifierOutputWithNoAttention:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.bit(pixel_values, output_hidden_states=output_hidden_states, return_dict=return_dict)
+
+ pooled_output = outputs.pooler_output if return_dict else outputs[1]
+
+ logits = self.classifier(pooled_output)
+
+ loss = None
+
+ if labels is not None:
+ loss = self.loss_function(labels, logits, self.config)
+
+ if not return_dict:
+ output = (logits,) + outputs[2:]
+ return (loss,) + output if loss is not None else output
+
+ return ImageClassifierOutputWithNoAttention(loss=loss, logits=logits, hidden_states=outputs.hidden_states)
+
+
+@auto_docstring(
+ custom_intro="""
+ BiT backbone, to be used with frameworks like DETR and MaskFormer.
+ """
+)
+class BitBackbone(BackboneMixin, BitPreTrainedModel):
+ has_attentions = False
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.bit = BitModel(config)
+ self.num_features = [config.embedding_size] + config.hidden_sizes
+
+ # initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @filter_output_hidden_states
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: Tensor,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> BackboneOutput:
+ r"""
+ Examples:
+
+ ```python
+ >>> from transformers import AutoImageProcessor, AutoBackbone
+ >>> import torch
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> processor = AutoImageProcessor.from_pretrained("google/bit-50")
+ >>> model = AutoBackbone.from_pretrained("google/bit-50")
+
+ >>> inputs = processor(image, return_tensors="pt")
+ >>> outputs = model(**inputs)
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+
+ outputs = self.bit(pixel_values, output_hidden_states=True, return_dict=True)
+
+ hidden_states = outputs.hidden_states
+
+ feature_maps = ()
+ for idx, stage in enumerate(self.stage_names):
+ if stage in self.out_features:
+ feature_maps += (hidden_states[idx],)
+
+ if not return_dict:
+ output = (feature_maps,)
+ if output_hidden_states:
+ output += (outputs.hidden_states,)
+ return output
+
+ return BackboneOutput(
+ feature_maps=feature_maps,
+ hidden_states=outputs.hidden_states if output_hidden_states else None,
+ attentions=None,
+ )
+
+
+__all__ = ["BitForImageClassification", "BitModel", "BitPreTrainedModel", "BitBackbone"]
diff --git a/third_party/transformers/src/transformers/models/blenderbot/__init__.py b/third_party/transformers/src/transformers/models/blenderbot/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..69beb03a4583078aaca1713ea66912f016af1a06
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/blenderbot/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_blenderbot import *
+ from .modeling_blenderbot import *
+ from .tokenization_blenderbot import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/blenderbot/configuration_blenderbot.py b/third_party/transformers/src/transformers/models/blenderbot/configuration_blenderbot.py
new file mode 100644
index 0000000000000000000000000000000000000000..26e3f1d7fec449645e33b4935a875007433fa14e
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/blenderbot/configuration_blenderbot.py
@@ -0,0 +1,81 @@
+# Copyright 2021 The Facebook, Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Blenderbot model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="facebook/blenderbot-3B")
+@strict
+class BlenderbotConfig(PreTrainedConfig):
+ r"""
+ encoder_no_repeat_ngram_size (`int`, *optional*, defaults to 3):
+ Number of ngrams to not be repeated in the encoder.
+
+ Example:
+
+ ```python
+ >>> from transformers import BlenderbotConfig, BlenderbotModel
+
+ >>> # Initializing a Blenderbot facebook/blenderbot-3B style configuration
+ >>> configuration = BlenderbotConfig()
+
+ >>> # Initializing a model (with random weights) from the facebook/blenderbot-3B style configuration
+ >>> model = BlenderbotModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "blenderbot"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ attribute_map = {
+ "num_attention_heads": "encoder_attention_heads",
+ "hidden_size": "d_model",
+ "num_hidden_layers": "encoder_layers",
+ }
+
+ vocab_size: int = 8008
+ max_position_embeddings: int = 128
+ encoder_layers: int = 2
+ encoder_ffn_dim: int = 10240
+ encoder_attention_heads: int = 32
+ decoder_layers: int = 24
+ decoder_ffn_dim: int = 10240
+ decoder_attention_heads: int = 32
+ encoder_layerdrop: float | int = 0.0
+ decoder_layerdrop: float | int = 0.0
+ use_cache: bool = True
+ is_encoder_decoder: bool = True
+ activation_function: str = "gelu"
+ d_model: int = 2560
+ dropout: float | int = 0.1
+ attention_dropout: float | int = 0.0
+ activation_dropout: float | int = 0.0
+ init_std: float = 0.02
+ decoder_start_token_id: int = 1
+ scale_embedding: bool = False
+ pad_token_id: int | None = 0
+ bos_token_id: int | None = 1
+ eos_token_id: int | list[int] | None = 2
+ encoder_no_repeat_ngram_size: int = 3
+ forced_eos_token_id: int | list[int] | None = 2
+ is_decoder: bool = False
+ tie_word_embeddings: bool = True
+
+
+__all__ = ["BlenderbotConfig"]
diff --git a/third_party/transformers/src/transformers/models/blenderbot/convert_blenderbot_original_pytorch_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/blenderbot/convert_blenderbot_original_pytorch_checkpoint_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..1ba9a3d570deed84d5154609b9b3ac0b4141f268
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/blenderbot/convert_blenderbot_original_pytorch_checkpoint_to_pytorch.py
@@ -0,0 +1,113 @@
+# Copyright 2020 The HuggingFace Inc. team.
+#
+# 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.
+"""Convert Blenderbot checkpoint."""
+
+import argparse
+
+import torch
+
+from transformers import BlenderbotConfig, BlenderbotForConditionalGeneration
+from transformers.utils import logging
+
+
+logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+PATTERNS = [
+ ["attention", "attn"],
+ ["encoder_attention", "encoder_attn"],
+ ["q_lin", "q_proj"],
+ ["k_lin", "k_proj"],
+ ["v_lin", "v_proj"],
+ ["out_lin", "out_proj"],
+ ["norm_embeddings", "layernorm_embedding"],
+ ["position_embeddings", "embed_positions"],
+ ["embeddings", "embed_tokens"],
+ ["ffn.lin", "fc"],
+]
+
+
+def rename_state_dict_key(k):
+ if k == "embeddings.weight":
+ return "shared.weight"
+
+ for parlai_name, hf_name in PATTERNS:
+ k = k.replace(parlai_name, hf_name)
+
+ if k.startswith("encoder"):
+ k = k.replace(".attn", ".self_attn")
+ k = k.replace("norm1", "self_attn_layer_norm")
+ k = k.replace("norm2", "final_layer_norm")
+ elif k.startswith("decoder"):
+ k = k.replace("norm1", "self_attn_layer_norm")
+ k = k.replace("norm2", "encoder_attn_layer_norm")
+ k = k.replace("norm3", "final_layer_norm")
+ return k
+
+
+def rename_layernorm_keys(sd):
+ keys = [
+ "model.encoder.layernorm_embedding.weight",
+ "model.encoder.layernorm_embedding.bias",
+ "model.decoder.layernorm_embedding.weight",
+ "model.decoder.layernorm_embedding.bias",
+ ]
+ for k in keys:
+ v = sd.pop(k)
+ new_k = k.replace("layernorm_embedding", "layer_norm")
+ assert new_k not in sd
+ sd[new_k] = v
+
+
+IGNORE_KEYS = ["START"]
+
+
+@torch.no_grad()
+def convert_parlai_checkpoint(checkpoint_path, pytorch_dump_folder_path, config_json_path):
+ """
+ Copy/paste/tweak model's weights to our BERT structure.
+ """
+ model = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+ sd = model["model"]
+ cfg = BlenderbotConfig.from_json_file(config_json_path)
+ m = BlenderbotForConditionalGeneration(cfg)
+ valid_keys = m.model.state_dict().keys()
+ failures = []
+ mapping = {}
+ for k, v in sd.items():
+ if k in IGNORE_KEYS:
+ continue
+
+ new_k = rename_state_dict_key(k)
+ if new_k not in valid_keys:
+ failures.append([k, new_k])
+ else:
+ mapping[new_k] = v
+ if cfg.normalize_before: # Blenderbot-3B checkpoints. Rename layernorm_embedding -> layer_norm
+ rename_layernorm_keys(sd)
+ m.model.load_state_dict(mapping, strict=True)
+ m.half()
+ m.save_pretrained(pytorch_dump_folder_path)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ # Required parameters
+ parser.add_argument("--src_path", type=str, help="like blenderbot-model.bin")
+ parser.add_argument("--save_dir", default="hf_blenderbot", type=str, help="Where to save converted model.")
+ parser.add_argument(
+ "--hf_config_json", default="blenderbot-3b-config.json", type=str, help="Path to config to use"
+ )
+ args = parser.parse_args()
+ convert_parlai_checkpoint(args.src_path, args.save_dir, args.hf_config_json)
diff --git a/third_party/transformers/src/transformers/models/blenderbot/modeling_blenderbot.py b/third_party/transformers/src/transformers/models/blenderbot/modeling_blenderbot.py
new file mode 100644
index 0000000000000000000000000000000000000000..99d5cb792234f2811e3a2f019366b71e831be716
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/blenderbot/modeling_blenderbot.py
@@ -0,0 +1,1022 @@
+# Copyright 2021 The Facebook, Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# 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 Blenderbot model."""
+
+import math
+from collections.abc import Callable
+
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_bidirectional_mask, create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPastAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions,
+ Seq2SeqLMOutput,
+ Seq2SeqModelOutput,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from .configuration_blenderbot import BlenderbotConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from transformers.models.bart.modeling_bart.shift_tokens_right
+def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
+ """
+ Shift input ids one token to the right.
+ """
+ shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+ shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
+ shifted_input_ids[:, 0] = decoder_start_token_id
+
+ if pad_token_id is None:
+ raise ValueError("self.model.config.pad_token_id has to be defined.")
+ # replace possible -100 values in labels by `pad_token_id`
+ shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+ return shifted_input_ids
+
+
+class BlenderbotLearnedPositionalEmbedding(nn.Embedding):
+ """
+ This module learns positional embeddings up to a fixed maximum size.
+ """
+
+ def __init__(self, num_embeddings: int, embedding_dim: int):
+ super().__init__(num_embeddings, embedding_dim)
+
+ def forward(
+ self, input_ids_shape: torch.Size, past_key_values_length: int = 0, position_ids: torch.Tensor | None = None
+ ):
+ """`input_ids_shape` is expected to be [bsz x seqlen]."""
+ if position_ids is None:
+ bsz, seq_len = input_ids_shape[:2]
+ position_ids = torch.arange(
+ past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device
+ )
+ return super().forward(position_ids)
+
+
+# Copied from transformers.models.bart.modeling_bart.BartScaledWordEmbedding with Bart->Blenderbot
+class BlenderbotScaledWordEmbedding(nn.Embedding):
+ """
+ This module overrides nn.Embeddings' forward by multiplying with embeddings scale.
+ """
+
+ def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: float | None = 1.0):
+ super().__init__(num_embeddings, embedding_dim, padding_idx)
+ self.embed_scale = embed_scale
+
+ def forward(self, input_ids: torch.Tensor):
+ return super().forward(input_ids) * self.embed_scale
+
+
+# Copied from transformers.models.bert.modeling_bert.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float | None = None,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ if scaling is None:
+ scaling = query.size(-1) ** -0.5
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.bart.modeling_bart.BartAttention with Bart->Blenderbot
+class BlenderbotAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ dropout: float = 0.0,
+ is_decoder: bool = False,
+ bias: bool = True,
+ is_causal: bool = False,
+ config: BlenderbotConfig | None = None,
+ layer_idx: int | None = None,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.dropout = dropout
+ self.head_dim = embed_dim // num_heads
+ self.config = config
+
+ if (self.head_dim * num_heads) != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+ f" and `num_heads`: {num_heads})."
+ )
+ self.scaling = self.head_dim**-0.5
+ self.is_decoder = is_decoder
+ self.is_causal = is_causal
+ self.layer_idx = layer_idx
+ if layer_idx is None and self.is_decoder:
+ logger.warning_once(
+ f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and "
+ "will lead to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class."
+ )
+
+ self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ key_value_states: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ # TODO: we need a refactor so that the different attention modules can get their specific kwargs
+ # ATM, we have mixed things encoder, decoder, and encoder-decoder attn
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Input shape: Batch x Time x Channel"""
+
+ # if key_value_states are provided this layer is used as a cross-attention layer
+ # for the decoder
+ is_cross_attention = key_value_states is not None
+
+ # determine input shapes
+ bsz, tgt_len = hidden_states.shape[:-1]
+ src_len = key_value_states.shape[1] if is_cross_attention else tgt_len
+
+ q_input_shape = (bsz, tgt_len, -1, self.head_dim)
+ kv_input_shape = (bsz, src_len, -1, self.head_dim)
+
+ # get query proj
+ query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2)
+
+ is_updated = False
+ if past_key_values is not None:
+ if isinstance(past_key_values, EncoderDecoderCache):
+ is_updated = past_key_values.is_updated.get(self.layer_idx)
+ if is_cross_attention:
+ # after the first generated id, we can subsequently re-use all key/value_states from cache
+ curr_past_key_values = past_key_values.cross_attention_cache
+ else:
+ curr_past_key_values = past_key_values.self_attention_cache
+ else:
+ curr_past_key_values = past_key_values
+
+ current_states = key_value_states if is_cross_attention else hidden_states
+ if is_cross_attention and past_key_values is not None and is_updated:
+ # reuse k,v, cross_attentions
+ key_states = curr_past_key_values.layers[self.layer_idx].keys
+ value_states = curr_past_key_values.layers[self.layer_idx].values
+ else:
+ key_states = self.k_proj(current_states)
+ value_states = self.v_proj(current_states)
+ key_states = key_states.view(*kv_input_shape).transpose(1, 2)
+ value_states = value_states.view(*kv_input_shape).transpose(1, 2)
+
+ if past_key_values is not None:
+ key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx)
+ # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+ if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):
+ past_key_values.is_updated[self.layer_idx] = True
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.mbart.modeling_mbart.MBartEncoderLayer with MBart->Blenderbot, MBART->BLENDERBOT
+class BlenderbotEncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: BlenderbotConfig):
+ super().__init__()
+ self.embed_dim = config.d_model
+
+ self.self_attn = BlenderbotAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.encoder_attention_heads,
+ dropout=config.attention_dropout,
+ config=config,
+ )
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.dropout = config.dropout
+ self.activation_fn = ACT2FN[config.activation_function]
+ self.activation_dropout = config.activation_dropout
+ self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
+ self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ """
+ residual = hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ if hidden_states.dtype == torch.float16:
+ clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+ hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+ return hidden_states
+
+
+# Copied from transformers.models.mbart.modeling_mbart.MBartDecoderLayer with MBart->Blenderbot, MBART->BLENDERBOT
+class BlenderbotDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: BlenderbotConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.embed_dim = config.d_model
+
+ self.self_attn = BlenderbotAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.decoder_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ is_causal=True,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.dropout = config.dropout
+ self.activation_fn = ACT2FN[config.activation_function]
+ self.activation_dropout = config.activation_dropout
+
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.encoder_attn = BlenderbotAttention(
+ self.embed_dim,
+ config.decoder_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
+ self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = True,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ encoder_hidden_states (`torch.FloatTensor`):
+ cross attention input to the layer of shape `(batch, seq_len, embed_dim)`
+ encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ past_key_values (`Cache`): cached past key and value projection states
+ """
+ residual = hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ # Cross-Attention Block
+ if encoder_hidden_states is not None:
+ residual = hidden_states
+ hidden_states = self.encoder_attn_layer_norm(hidden_states)
+
+ hidden_states, _ = self.encoder_attn(
+ hidden_states=hidden_states,
+ key_value_states=encoder_hidden_states,
+ attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+@auto_docstring
+class BlenderbotPreTrainedModel(PreTrainedModel):
+ config: BlenderbotConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+ _can_compile_fullgraph = True
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, BlenderbotForConditionalGeneration):
+ init.zeros_(module.final_logits_bias)
+
+ @property
+ def dummy_inputs(self):
+ pad_token = self.config.pad_token_id
+ input_ids = torch.tensor([[0, 6, 10, 4, 2], [0, 8, 12, 2, pad_token]], device=self.device)
+ dummy_inputs = {
+ "attention_mask": input_ids.ne(pad_token),
+ "input_ids": input_ids,
+ "decoder_input_ids": input_ids,
+ }
+ return dummy_inputs
+
+
+class BlenderbotEncoder(BlenderbotPreTrainedModel):
+ """
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
+ [`BlenderbotEncoderLayer`].
+
+ Args:
+ config: BlenderbotConfig
+ embed_tokens (nn.Embedding): output embedding
+ """
+
+ _can_record_outputs = {
+ "hidden_states": BlenderbotEncoderLayer,
+ "attentions": BlenderbotAttention,
+ }
+
+ def __init__(self, config: BlenderbotConfig):
+ super().__init__(config)
+
+ self.dropout = config.dropout
+ self.layerdrop = config.encoder_layerdrop
+
+ embed_dim = config.d_model
+ self.padding_idx = config.pad_token_id
+ self.max_source_positions = config.max_position_embeddings
+ embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+
+ self.embed_tokens = BlenderbotScaledWordEmbedding(
+ config.vocab_size, embed_dim, self.padding_idx, embed_scale=embed_scale
+ )
+
+ self.embed_positions = BlenderbotLearnedPositionalEmbedding(
+ config.max_position_embeddings,
+ embed_dim,
+ )
+ self.layers = nn.ModuleList([BlenderbotEncoderLayer(config) for _ in range(config.encoder_layers)])
+ self.layer_norm = nn.LayerNorm(config.d_model)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutput:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ input_shape = inputs_embeds.size()[:-1]
+ embed_pos = self.embed_positions(input_shape)
+
+ hidden_states = inputs_embeds + embed_pos
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ )
+ for idx, encoder_layer in enumerate(self.layers):
+ # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+ to_drop = False
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop: # skip the layer
+ to_drop = True
+
+ if not to_drop:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask,
+ **kwargs,
+ )
+
+ # add final layer norm
+ hidden_states = self.layer_norm(hidden_states)
+
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+class BlenderbotDecoder(BlenderbotPreTrainedModel):
+ """
+ Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`BlenderbotDecoderLayer`]
+
+ Args:
+ config: BlenderbotConfig
+ embed_tokens (nn.Embedding): output embedding
+ """
+
+ _can_record_outputs = {
+ "hidden_states": BlenderbotDecoderLayer,
+ "attentions": OutputRecorder(BlenderbotAttention, index=1, layer_name="self_attn"),
+ "cross_attentions": OutputRecorder(BlenderbotAttention, index=1, layer_name="encoder_attn"),
+ }
+
+ def __init__(self, config: BlenderbotConfig):
+ super().__init__(config)
+ self.dropout = config.dropout
+ self.layerdrop = config.decoder_layerdrop
+ self.padding_idx = config.pad_token_id
+ self.max_target_positions = config.max_position_embeddings
+ embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
+
+ self.embed_tokens = BlenderbotScaledWordEmbedding(
+ config.vocab_size, config.d_model, self.padding_idx, embed_scale=embed_scale
+ )
+
+ self.embed_positions = BlenderbotLearnedPositionalEmbedding(
+ config.max_position_embeddings,
+ config.d_model,
+ )
+ self.layers = nn.ModuleList(
+ [BlenderbotDecoderLayer(config, layer_idx=i) for i in range(config.decoder_layers)]
+ )
+ self.layer_norm = nn.LayerNorm(config.d_model)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPastAndCrossAttentions:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ # initialize `past_key_values`
+ if use_cache and past_key_values is None:
+ past_key_values = (
+ EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+ if encoder_hidden_states is not None or self.config.is_encoder_decoder
+ else DynamicCache(config=self.config)
+ )
+
+ batch_size, seq_length = inputs_embeds.size()[:-1]
+ past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(seq_length, device=inputs_embeds.device) + past_key_values_length
+
+ if attention_mask is None and not is_torchdynamo_compiling():
+ # required mask seq length can be calculated via length of past cache
+ mask_seq_length = past_key_values_length + seq_length
+ attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)
+
+ self_attn_cache = (
+ past_key_values.self_attention_cache
+ if isinstance(past_key_values, EncoderDecoderCache)
+ else past_key_values
+ )
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=self_attn_cache,
+ )
+ encoder_attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=encoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ )
+
+ # embed positions
+ position_ids = self.embed_positions(
+ (batch_size, seq_length), past_key_values_length, position_ids=position_ids
+ )
+
+ hidden_states = inputs_embeds + position_ids
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ for idx, decoder_layer in enumerate(self.layers):
+ # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop:
+ continue
+
+ layer_outputs = decoder_layer(
+ hidden_states,
+ causal_mask,
+ encoder_hidden_states, # as a positional argument for gradient checkpointing
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+ hidden_states = layer_outputs[0] if isinstance(layer_outputs, tuple) else layer_outputs
+
+ # add final layer norm
+ hidden_states = self.layer_norm(hidden_states)
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class BlenderbotModel(BlenderbotPreTrainedModel):
+ _tied_weights_keys = {
+ "encoder.embed_tokens.weight": "shared.weight",
+ "decoder.embed_tokens.weight": "shared.weight",
+ }
+
+ def __init__(self, config: BlenderbotConfig):
+ super().__init__(config)
+
+ padding_idx, vocab_size = config.pad_token_id, config.vocab_size
+ embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
+ self.shared = BlenderbotScaledWordEmbedding(vocab_size, config.d_model, padding_idx, embed_scale=embed_scale)
+ self.encoder = BlenderbotEncoder(config)
+ self.decoder = BlenderbotDecoder(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.shared
+
+ def set_input_embeddings(self, value):
+ self.shared = value
+ self.encoder.embed_tokens = self.shared
+ self.decoder.embed_tokens = self.shared
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: BaseModelOutput | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ decoder_inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Seq2SeqModelOutput:
+ r"""
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Indices of decoder input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+ Blenderbot uses the `bos_token_id` as the starting token for `decoder_input_ids` generation. If
+ `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
+ `past_key_values`).
+ decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+ be used by default.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, BlenderbotModel
+
+ >>> model = BlenderbotModel.from_pretrained("facebook/blenderbot-400M-distill")
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-400M-distill")
+
+ >>> inputs = tokenizer("Studies have been shown that owning a dog is good for you", return_tensors="pt")
+ >>> decoder_input_ids = tokenizer("Studies show that", return_tensors="pt").input_ids # Batch size 1
+ >>> outputs = model(input_ids=inputs.input_ids, decoder_input_ids=decoder_input_ids)
+
+ >>> last_hidden_states = outputs.last_hidden_state
+ >>> list(last_hidden_states.shape)
+ [1, 6, 1280]
+ ```"""
+
+ if encoder_outputs is None:
+ encoder_outputs: BaseModelOutput = self.encoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+ elif not isinstance(encoder_outputs, BaseModelOutput):
+ encoder_outputs = BaseModelOutput(
+ last_hidden_state=encoder_outputs[0],
+ hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
+ attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
+ )
+
+ decoder_outputs: BaseModelOutputWithPastAndCrossAttentions = self.decoder(
+ input_ids=decoder_input_ids,
+ attention_mask=decoder_attention_mask,
+ encoder_hidden_states=encoder_outputs[0],
+ encoder_attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=decoder_inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ return Seq2SeqModelOutput(
+ last_hidden_state=decoder_outputs.last_hidden_state,
+ past_key_values=decoder_outputs.past_key_values,
+ decoder_hidden_states=decoder_outputs.hidden_states,
+ decoder_attentions=decoder_outputs.attentions,
+ cross_attentions=decoder_outputs.cross_attentions,
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+ encoder_hidden_states=encoder_outputs.hidden_states,
+ encoder_attentions=encoder_outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The Blenderbot Model with a language modeling head. Can be used for summarization.
+ """
+)
+class BlenderbotForConditionalGeneration(BlenderbotPreTrainedModel, GenerationMixin):
+ base_model_prefix = "model"
+ _keys_to_ignore_on_load_missing = ["final_logits_bias"]
+ _tied_weights_keys = {
+ "lm_head.weight": "model.shared.weight",
+ }
+
+ def __init__(self, config: BlenderbotConfig):
+ super().__init__(config)
+ self.model = BlenderbotModel(config)
+ self.register_buffer("final_logits_bias", torch.zeros((1, self.model.shared.num_embeddings)))
+ self.lm_head = nn.Linear(config.d_model, self.model.shared.num_embeddings, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def resize_token_embeddings(
+ self, new_num_tokens: int, pad_to_multiple_of: int | None = None, mean_resizing: bool = True
+ ) -> nn.Embedding:
+ new_embeddings = super().resize_token_embeddings(new_num_tokens, pad_to_multiple_of, mean_resizing)
+ self._resize_final_logits_bias(new_embeddings.weight.shape[0])
+ return new_embeddings
+
+ def _resize_final_logits_bias(self, new_num_tokens: int) -> None:
+ old_num_tokens = self.final_logits_bias.shape[-1]
+ if new_num_tokens <= old_num_tokens:
+ new_bias = self.final_logits_bias[:, :new_num_tokens]
+ else:
+ extra_bias = torch.zeros((1, new_num_tokens - old_num_tokens), device=self.final_logits_bias.device)
+ new_bias = torch.cat([self.final_logits_bias, extra_bias], dim=1)
+ self.register_buffer("final_logits_bias", new_bias)
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: BaseModelOutput | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ decoder_inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Seq2SeqLMOutput:
+ r"""
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Indices of decoder input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+ Blenderbot uses the `bos_token_id` as the starting token for `decoder_input_ids` generation. If
+ `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
+ `past_key_values`).
+ decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+ be used by default.
+ 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]`.
+
+ Example conversation:
+
+ ```python
+ >>> from transformers import AutoTokenizer, BlenderbotForConditionalGeneration
+
+ >>> mname = "facebook/blenderbot-400M-distill"
+ >>> model = BlenderbotForConditionalGeneration.from_pretrained(mname)
+ >>> tokenizer = AutoTokenizer.from_pretrained(mname)
+ >>> UTTERANCE = "My friends are cool but they eat too many carbs."
+ >>> print("Human: ", UTTERANCE)
+ Human: My friends are cool but they eat too many carbs.
+
+ >>> inputs = tokenizer([UTTERANCE], return_tensors="pt")
+ >>> reply_ids = model.generate(**inputs)
+ >>> print("Bot: ", tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0])
+ Bot: That's unfortunate. Are they trying to lose weight or are they just trying to be healthier?
+
+ >>> REPLY = "I'm not sure"
+ >>> print("Human: ", REPLY)
+ Human: I'm not sure
+
+ >>> NEXT_UTTERANCE = (
+ ... "My friends are cool but they eat too many carbs. That's unfortunate. "
+ ... "Are they trying to lose weight or are they just trying to be healthier? "
+ ... " I'm not sure."
+ ... )
+ >>> inputs = tokenizer([NEXT_UTTERANCE], return_tensors="pt")
+ >>> next_reply_ids = model.generate(**inputs)
+ >>> print("Bot: ", tokenizer.batch_decode(next_reply_ids, skip_special_tokens=True)[0])
+ Bot: I see. Well, it's good that they're trying to change their eating habits.
+ ```
+ """
+ if labels is not None:
+ if use_cache:
+ logger.warning("The `use_cache` argument is changed to `False` since `labels` is provided.")
+ use_cache = False
+ if decoder_input_ids is None and decoder_inputs_embeds is None:
+ decoder_input_ids = shift_tokens_right(
+ labels, self.config.pad_token_id, self.config.decoder_start_token_id
+ )
+
+ outputs: Seq2SeqModelOutput = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ decoder_input_ids=decoder_input_ids,
+ encoder_outputs=encoder_outputs,
+ decoder_attention_mask=decoder_attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ decoder_inputs_embeds=decoder_inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ lm_logits = self.lm_head(outputs[0])
+ lm_logits = lm_logits + self.final_logits_bias.to(lm_logits.device)
+
+ masked_lm_loss = None
+ if labels is not None:
+ labels = labels.to(lm_logits.device)
+ loss_fct = CrossEntropyLoss()
+ masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+ return Seq2SeqLMOutput(
+ loss=masked_lm_loss,
+ logits=lm_logits,
+ past_key_values=outputs.past_key_values,
+ decoder_hidden_states=outputs.decoder_hidden_states,
+ decoder_attentions=outputs.decoder_attentions,
+ cross_attentions=outputs.cross_attentions,
+ encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+ encoder_hidden_states=outputs.encoder_hidden_states,
+ encoder_attentions=outputs.encoder_attentions,
+ )
+
+
+# Copied from transformers.models.bart.modeling_bart.BartDecoderWrapper with Bart->Blenderbot
+class BlenderbotDecoderWrapper(BlenderbotPreTrainedModel):
+ """
+ This wrapper class is a helper class to correctly load pretrained checkpoints when the causal language model is
+ used in combination with the [`EncoderDecoderModel`] framework.
+ """
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.decoder = BlenderbotDecoder(config)
+ self.post_init()
+
+ def forward(self, *args, **kwargs):
+ return self.decoder(*args, **kwargs)
+
+
+# Copied from transformers.models.bart.modeling_bart.BartForCausalLM with Bart->Blenderbot, facebook/bart-base->facebook/blenderbot-400M-distill
+class BlenderbotForCausalLM(BlenderbotPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {
+ "lm_head.weight": "model.decoder.embed_tokens.weight",
+ }
+
+ def __init__(self, config):
+ config.is_decoder = True
+ config.is_encoder_decoder = False
+ super().__init__(config)
+ self.model = BlenderbotDecoderWrapper(config)
+
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.decoder.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.model.decoder.embed_tokens = value
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithCrossAttentions:
+ r"""
+ 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]`.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, BlenderbotForCausalLM
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot-400M-distill")
+ >>> model = BlenderbotForCausalLM.from_pretrained("facebook/blenderbot-400M-distill")
+ >>> assert model.config.is_decoder, f"{model.__class__} has to be configured as a decoder."
+ >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+ >>> outputs = model(**inputs)
+
+ >>> logits = outputs.logits
+ >>> expected_shape = [1, inputs.input_ids.shape[-1], model.config.vocab_size]
+ >>> list(logits.shape) == expected_shape
+ True
+ ```"""
+
+ outputs: BaseModelOutputWithPastAndCrossAttentions = self.model.decoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ # Only compute necessary logits
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ labels = labels.to(logits.device)
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+ return CausalLMOutputWithCrossAttentions(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+__all__ = [
+ "BlenderbotForCausalLM",
+ "BlenderbotForConditionalGeneration",
+ "BlenderbotModel",
+ "BlenderbotPreTrainedModel",
+]
diff --git a/third_party/transformers/src/transformers/models/blenderbot/tokenization_blenderbot.py b/third_party/transformers/src/transformers/models/blenderbot/tokenization_blenderbot.py
new file mode 100644
index 0000000000000000000000000000000000000000..6d1951cd3783226866d880c63015b74558b265be
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/blenderbot/tokenization_blenderbot.py
@@ -0,0 +1,175 @@
+# Copyright 2021 The Facebook Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Tokenization class for Blenderbot."""
+
+from tokenizers import Tokenizer, decoders, pre_tokenizers
+from tokenizers.models import BPE
+
+from ...tokenization_utils_base import AddedToken
+from ...tokenization_utils_tokenizers import TokenizersBackend
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+VOCAB_FILES_NAMES = {
+ "vocab_file": "vocab.json",
+ "merges_file": "merges.txt",
+ "tokenizer_config_file": "tokenizer_config.json",
+}
+
+
+class BlenderbotTokenizer(TokenizersBackend):
+ """
+ Construct a "fast" Blenderbot tokenizer (backed by HuggingFace's *tokenizers* library), derived from the GPT-2
+ tokenizer, using byte-level Byte-Pair-Encoding.
+
+ This tokenizer has been trained to treat spaces like parts of the tokens (a bit like sentencepiece) so a word will
+ be encoded differently whether it is at the beginning of the sentence (without space) or not:
+
+ ```python
+ >>> from transformers import BlenderbotTokenizerFast
+
+ >>> tokenizer = BlenderbotTokenizerFast.from_pretrained("facebook/blenderbot-3B")
+ >>> tokenizer("Hello world")["input_ids"]
+ [6950, 1085, 2]
+
+ >>> tokenizer(" Hello world")["input_ids"]
+ [6950, 1085, 2]
+ ```
+
+ You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer or when you
+ call it on some text, but since the model was not pretrained this way, it might yield a decrease in performance.
+
+
+
+ When used with `is_split_into_words=True`, this tokenizer needs to be instantiated with `add_prefix_space=True`.
+
+
+
+ This tokenizer inherits from [`PreTrainedTokenizerFast`] which contains most of the main methods. Users should
+ refer to this superclass for more information regarding those methods.
+
+ Args:
+ bos_token (`str`, *optional*, defaults to `""`):
+ The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
+
+
+
+ When building a sequence using special tokens, this is not the token that is used for the beginning of
+ sequence. The token used is the `cls_token`.
+
+
+
+ eos_token (`str`, *optional*, defaults to `""`):
+ The end of sequence token.
+
+
+
+ When building a sequence using special tokens, this is not the token that is used for the end of sequence.
+ The token used is the `sep_token`.
+
+
+
+ sep_token (`str`, *optional*, defaults to `""`):
+ The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+ sequence classification or for a text and a question for question answering. It is also used as the last
+ token of a sequence built with special tokens.
+ cls_token (`str`, *optional*, defaults to `""`):
+ The classifier token which is used when doing sequence classification (classification of the whole sequence
+ instead of per-token classification). It is the first token of the sequence when built with special tokens.
+ unk_token (`str`, *optional*, defaults to `""`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ pad_token (`str`, *optional*, defaults to `""`):
+ The token used for padding, for example when batching sequences of different lengths.
+ mask_token (`str`, *optional*, defaults to `""`):
+ The token used for masking values. This is the token used when training this model with masked language
+ modeling. This is the token which the model will try to predict.
+ add_prefix_space (`bool`, *optional*, defaults to `True`):
+ Whether or not to add an initial space to the input. This allows to treat the leading word just as any
+ other word. (Blenderbot tokenizer detect beginning of words by the preceding space).
+ vocab (`str` or `dict[str, int]`, *optional*):
+ Custom vocabulary dictionary. If not provided, vocabulary is loaded from `vocab_file`.
+ merges (`str` or `list[str]`, *optional*):
+ Custom merges list. If not provided, merges are loaded from `merges_file`.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ model_input_names = ["input_ids", "attention_mask"]
+ model = BPE
+
+ def __init__(
+ self,
+ bos_token="",
+ eos_token="",
+ sep_token="",
+ cls_token="",
+ unk_token="",
+ pad_token="",
+ mask_token="",
+ add_prefix_space=True,
+ vocab=None,
+ merges=None,
+ **kwargs,
+ ):
+ self.add_prefix_space = add_prefix_space
+ mask_token = (
+ AddedToken(mask_token, lstrip=True, rstrip=False, normalized=False)
+ if isinstance(mask_token, str)
+ else mask_token
+ )
+
+ # Initialize vocab and merges; when not provided fall back to minimal vocab
+ self._vocab = (
+ vocab
+ if vocab is not None
+ else {
+ str(bos_token): 0,
+ str(pad_token): 1,
+ str(eos_token): 2,
+ str(unk_token): 3,
+ str(mask_token): 4,
+ }
+ )
+
+ self._merges = merges or []
+ self._tokenizer = Tokenizer(
+ BPE(
+ vocab=self._vocab,
+ merges=self._merges,
+ dropout=None,
+ continuing_subword_prefix="",
+ end_of_word_suffix="",
+ fuse_unk=False,
+ )
+ )
+
+ self._tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=add_prefix_space)
+ self._tokenizer.decoder = decoders.ByteLevel()
+ super().__init__(
+ bos_token=bos_token,
+ eos_token=eos_token,
+ sep_token=sep_token,
+ cls_token=cls_token,
+ unk_token=unk_token,
+ pad_token=pad_token,
+ mask_token=mask_token,
+ add_prefix_space=add_prefix_space,
+ **kwargs,
+ )
+
+
+__all__ = ["BlenderbotTokenizer"]
diff --git a/third_party/transformers/src/transformers/models/blenderbot_small/__init__.py b/third_party/transformers/src/transformers/models/blenderbot_small/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..7f08df82e757b33886ee65d8b6ad050ea663fe78
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/blenderbot_small/__init__.py
@@ -0,0 +1,29 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_blenderbot_small import *
+ from .modeling_blenderbot_small import *
+ from .tokenization_blenderbot_small import *
+ from .tokenization_blenderbot_small_fast import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/blenderbot_small/configuration_blenderbot_small.py b/third_party/transformers/src/transformers/models/blenderbot_small/configuration_blenderbot_small.py
new file mode 100644
index 0000000000000000000000000000000000000000..f70d2a36fb224b2d8929e1eb66b1788ac4f566e9
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/blenderbot_small/configuration_blenderbot_small.py
@@ -0,0 +1,77 @@
+# Copyright 2021 The Facebook, Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""BlenderbotSmall model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="facebook/blenderbot_small-90M")
+@strict
+class BlenderbotSmallConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import BlenderbotSmallConfig, BlenderbotSmallModel
+
+ >>> # Initializing a BlenderbotSmall facebook/blenderbot_small-90M style configuration
+ >>> configuration = BlenderbotSmallConfig()
+
+ >>> # Initializing a model (with random weights) from the facebook/blenderbot_small-90M style configuration
+ >>> model = BlenderbotSmallModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "blenderbot-small"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ attribute_map = {
+ "num_attention_heads": "encoder_attention_heads",
+ "hidden_size": "d_model",
+ "num_hidden_layers": "encoder_layers",
+ }
+
+ vocab_size: int = 50265
+ max_position_embeddings: int = 512
+ encoder_layers: int = 8
+ encoder_ffn_dim: int = 2048
+ encoder_attention_heads: int = 16
+ decoder_layers: int = 8
+ decoder_ffn_dim: int = 2048
+ decoder_attention_heads: int = 16
+ encoder_layerdrop: float | int = 0.0
+ decoder_layerdrop: float | int = 0.0
+ use_cache: bool = True
+ is_encoder_decoder: bool = True
+ activation_function: str = "gelu"
+ d_model: int = 512
+ dropout: float | int = 0.1
+ attention_dropout: float | int = 0.0
+ activation_dropout: float | int = 0.0
+ init_std: float = 0.02
+ decoder_start_token_id: int | None = 1
+ scale_embedding: bool = False
+ pad_token_id: int | None = 0
+ bos_token_id: int | None = 1
+ eos_token_id: int | list[int] | None = 2
+ forced_eos_token_id: int | list[int] | None = 2
+ is_decoder: bool = False
+ tie_word_embeddings: bool = True
+
+
+__all__ = ["BlenderbotSmallConfig"]
diff --git a/third_party/transformers/src/transformers/models/blenderbot_small/modeling_blenderbot_small.py b/third_party/transformers/src/transformers/models/blenderbot_small/modeling_blenderbot_small.py
new file mode 100644
index 0000000000000000000000000000000000000000..89932e6f282cfb3612301dcb48868ebb74309548
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/blenderbot_small/modeling_blenderbot_small.py
@@ -0,0 +1,986 @@
+# Copyright 2021 The Facebook, Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# 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 BlenderbotSmall model."""
+
+import math
+from collections.abc import Callable
+
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_bidirectional_mask, create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPastAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions,
+ Seq2SeqLMOutput,
+ Seq2SeqModelOutput,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from .configuration_blenderbot_small import BlenderbotSmallConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from transformers.models.bart.modeling_bart.shift_tokens_right
+def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
+ """
+ Shift input ids one token to the right.
+ """
+ shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+ shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
+ shifted_input_ids[:, 0] = decoder_start_token_id
+
+ if pad_token_id is None:
+ raise ValueError("self.model.config.pad_token_id has to be defined.")
+ # replace possible -100 values in labels by `pad_token_id`
+ shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+ return shifted_input_ids
+
+
+# Copied from transformers.models.blenderbot.modeling_blenderbot.BlenderbotLearnedPositionalEmbedding with Blenderbot->BlenderbotSmall
+class BlenderbotSmallLearnedPositionalEmbedding(nn.Embedding):
+ """
+ This module learns positional embeddings up to a fixed maximum size.
+ """
+
+ def __init__(self, num_embeddings: int, embedding_dim: int):
+ super().__init__(num_embeddings, embedding_dim)
+
+ def forward(
+ self, input_ids_shape: torch.Size, past_key_values_length: int = 0, position_ids: torch.Tensor | None = None
+ ):
+ """`input_ids_shape` is expected to be [bsz x seqlen]."""
+ if position_ids is None:
+ bsz, seq_len = input_ids_shape[:2]
+ position_ids = torch.arange(
+ past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device
+ )
+ return super().forward(position_ids)
+
+
+# Copied from transformers.models.bert.modeling_bert.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float | None = None,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ if scaling is None:
+ scaling = query.size(-1) ** -0.5
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.bart.modeling_bart.BartAttention with Bart->BlenderbotSmall
+class BlenderbotSmallAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ dropout: float = 0.0,
+ is_decoder: bool = False,
+ bias: bool = True,
+ is_causal: bool = False,
+ config: BlenderbotSmallConfig | None = None,
+ layer_idx: int | None = None,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.dropout = dropout
+ self.head_dim = embed_dim // num_heads
+ self.config = config
+
+ if (self.head_dim * num_heads) != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+ f" and `num_heads`: {num_heads})."
+ )
+ self.scaling = self.head_dim**-0.5
+ self.is_decoder = is_decoder
+ self.is_causal = is_causal
+ self.layer_idx = layer_idx
+ if layer_idx is None and self.is_decoder:
+ logger.warning_once(
+ f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and "
+ "will lead to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class."
+ )
+
+ self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ key_value_states: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ # TODO: we need a refactor so that the different attention modules can get their specific kwargs
+ # ATM, we have mixed things encoder, decoder, and encoder-decoder attn
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Input shape: Batch x Time x Channel"""
+
+ # if key_value_states are provided this layer is used as a cross-attention layer
+ # for the decoder
+ is_cross_attention = key_value_states is not None
+
+ # determine input shapes
+ bsz, tgt_len = hidden_states.shape[:-1]
+ src_len = key_value_states.shape[1] if is_cross_attention else tgt_len
+
+ q_input_shape = (bsz, tgt_len, -1, self.head_dim)
+ kv_input_shape = (bsz, src_len, -1, self.head_dim)
+
+ # get query proj
+ query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2)
+
+ is_updated = False
+ if past_key_values is not None:
+ if isinstance(past_key_values, EncoderDecoderCache):
+ is_updated = past_key_values.is_updated.get(self.layer_idx)
+ if is_cross_attention:
+ # after the first generated id, we can subsequently re-use all key/value_states from cache
+ curr_past_key_values = past_key_values.cross_attention_cache
+ else:
+ curr_past_key_values = past_key_values.self_attention_cache
+ else:
+ curr_past_key_values = past_key_values
+
+ current_states = key_value_states if is_cross_attention else hidden_states
+ if is_cross_attention and past_key_values is not None and is_updated:
+ # reuse k,v, cross_attentions
+ key_states = curr_past_key_values.layers[self.layer_idx].keys
+ value_states = curr_past_key_values.layers[self.layer_idx].values
+ else:
+ key_states = self.k_proj(current_states)
+ value_states = self.v_proj(current_states)
+ key_states = key_states.view(*kv_input_shape).transpose(1, 2)
+ value_states = value_states.view(*kv_input_shape).transpose(1, 2)
+
+ if past_key_values is not None:
+ key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx)
+ # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+ if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):
+ past_key_values.is_updated[self.layer_idx] = True
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.bart.modeling_bart.BartEncoderLayer with Bart->BlenderbotSmall, BART->BLENDERBOT_SMALL
+class BlenderbotSmallEncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: BlenderbotSmallConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.embed_dim = config.d_model
+
+ self.self_attn = BlenderbotSmallAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.encoder_attention_heads,
+ dropout=config.attention_dropout,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.dropout = config.dropout
+ self.activation_fn = ACT2FN[config.activation_function]
+ self.activation_dropout = config.activation_dropout
+ self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)
+ self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor,
+ attention_mask: torch.FloatTensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states, _ = self.self_attn(
+ hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ residual = hidden_states
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+
+ if hidden_states.dtype == torch.float16 and not torch.isfinite(hidden_states).all():
+ clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+ hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)
+
+ return hidden_states
+
+
+# Copied from transformers.models.bart.modeling_bart.BartDecoderLayer with Bart->BlenderbotSmall, BART->BLENDERBOT_SMALL
+class BlenderbotSmallDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: BlenderbotSmallConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.embed_dim = config.d_model
+
+ self.self_attn = BlenderbotSmallAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.decoder_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ is_causal=True,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.dropout = config.dropout
+ self.activation_fn = ACT2FN[config.activation_function]
+ self.activation_dropout = config.activation_dropout
+
+ self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.encoder_attn = BlenderbotSmallAttention(
+ self.embed_dim,
+ config.decoder_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=True,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+ self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
+ self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
+ self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = True,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ # Cross-Attention Block
+ if encoder_hidden_states is not None:
+ residual = hidden_states
+
+ hidden_states, _ = self.encoder_attn(
+ hidden_states,
+ key_value_states=encoder_hidden_states,
+ attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states = self.encoder_attn_layer_norm(hidden_states)
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
+ hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+
+ return hidden_states
+
+
+@auto_docstring
+class BlenderbotSmallPreTrainedModel(PreTrainedModel):
+ config: BlenderbotSmallConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+ _can_compile_fullgraph = True
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, BlenderbotSmallForConditionalGeneration):
+ init.zeros_(module.final_logits_bias)
+
+ @property
+ def dummy_inputs(self):
+ pad_token = self.config.pad_token_id
+ input_ids = torch.tensor([[0, 6, 10, 4, 2], [0, 8, 12, 2, pad_token]], device=self.device)
+ dummy_inputs = {
+ "attention_mask": input_ids.ne(pad_token),
+ "input_ids": input_ids,
+ "decoder_input_ids": input_ids,
+ }
+ return dummy_inputs
+
+
+class BlenderbotSmallEncoder(BlenderbotSmallPreTrainedModel):
+ """
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
+ [`BlenderbotSmallEncoderLayer`].
+
+ Args:
+ config: BlenderbotSmallConfig
+ embed_tokens (nn.Embedding): output embedding
+ """
+
+ _can_record_outputs = {
+ "hidden_states": BlenderbotSmallEncoderLayer,
+ "attentions": BlenderbotSmallAttention,
+ }
+
+ def __init__(self, config: BlenderbotSmallConfig):
+ super().__init__(config)
+
+ self.dropout = config.dropout
+ self.layerdrop = config.encoder_layerdrop
+
+ embed_dim = config.d_model
+ self.padding_idx = config.pad_token_id
+ self.max_source_positions = config.max_position_embeddings
+ self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, embed_dim, self.padding_idx)
+
+ self.embed_positions = BlenderbotSmallLearnedPositionalEmbedding(
+ config.max_position_embeddings,
+ embed_dim,
+ )
+ self.layers = nn.ModuleList([BlenderbotSmallEncoderLayer(config) for _ in range(config.encoder_layers)])
+ self.layernorm_embedding = nn.LayerNorm(embed_dim)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutput:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale
+
+ input_shape = inputs_embeds.size()[:-1]
+ embed_pos = self.embed_positions(input_shape)
+
+ hidden_states = inputs_embeds + embed_pos
+ hidden_states = self.layernorm_embedding(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ )
+
+ for idx, encoder_layer in enumerate(self.layers):
+ # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+ to_drop = False
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop: # skip the layer
+ to_drop = True
+
+ if not to_drop:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask,
+ **kwargs,
+ )
+
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+class BlenderbotSmallDecoder(BlenderbotSmallPreTrainedModel):
+ """
+ Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`BlenderbotSmallDecoderLayer`]
+
+ Args:
+ config: BlenderbotSmallConfig
+ embed_tokens (nn.Embedding): output embedding
+ """
+
+ _can_record_outputs = {
+ "hidden_states": BlenderbotSmallDecoderLayer,
+ "attentions": OutputRecorder(BlenderbotSmallAttention, index=1, layer_name="self_attn"),
+ "cross_attentions": OutputRecorder(BlenderbotSmallAttention, index=1, layer_name="encoder_attn"),
+ }
+
+ def __init__(self, config: BlenderbotSmallConfig):
+ super().__init__(config)
+ self.dropout = config.dropout
+ self.layerdrop = config.decoder_layerdrop
+ self.padding_idx = config.pad_token_id
+ self.max_target_positions = config.max_position_embeddings
+ self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model, self.padding_idx)
+
+ self.embed_positions = BlenderbotSmallLearnedPositionalEmbedding(
+ config.max_position_embeddings,
+ config.d_model,
+ )
+ self.layers = nn.ModuleList(
+ [BlenderbotSmallDecoderLayer(config, layer_idx=i) for i in range(config.decoder_layers)]
+ )
+ self.layernorm_embedding = nn.LayerNorm(config.d_model)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids=None,
+ attention_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ past_key_values=None,
+ inputs_embeds=None,
+ use_cache=None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPastAndCrossAttentions:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ # initialize `past_key_values`
+ if use_cache and past_key_values is None:
+ past_key_values = (
+ EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+ if encoder_hidden_states is not None or self.config.is_encoder_decoder
+ else DynamicCache(config=self.config)
+ )
+
+ batch_size, seq_length = inputs_embeds.size()[:-1]
+ past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(seq_length, device=inputs_embeds.device) + past_key_values_length
+
+ if attention_mask is None and not is_torchdynamo_compiling():
+ # required mask seq length can be calculated via length of past cache
+ mask_seq_length = past_key_values_length + seq_length
+ attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)
+
+ self_attn_cache = (
+ past_key_values.self_attention_cache
+ if isinstance(past_key_values, EncoderDecoderCache)
+ else past_key_values
+ )
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=self_attn_cache,
+ )
+ encoder_attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=encoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ )
+
+ # embed positions
+ position_ids = self.embed_positions(
+ (batch_size, seq_length), past_key_values_length, position_ids=position_ids
+ )
+
+ # BlenderbotSmall applies layer norm on hidden_states
+ inputs_embeds = self.layernorm_embedding(inputs_embeds)
+ hidden_states = inputs_embeds + position_ids
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+
+ for idx, decoder_layer in enumerate(self.layers):
+ # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)
+ if self.training:
+ dropout_probability = torch.rand([])
+ if dropout_probability < self.layerdrop:
+ continue
+
+ hidden_states = decoder_layer(
+ hidden_states,
+ causal_mask,
+ encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class BlenderbotSmallModel(BlenderbotSmallPreTrainedModel):
+ _tied_weights_keys = {
+ "encoder.embed_tokens.weight": "shared.weight",
+ "decoder.embed_tokens.weight": "shared.weight",
+ }
+
+ def __init__(self, config: BlenderbotSmallConfig):
+ super().__init__(config)
+
+ padding_idx, vocab_size = config.pad_token_id, config.vocab_size
+ self.shared = nn.Embedding(vocab_size, config.d_model, padding_idx)
+
+ self.encoder = BlenderbotSmallEncoder(config)
+ self.decoder = BlenderbotSmallDecoder(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.shared
+
+ def set_input_embeddings(self, value):
+ self.shared = value
+ self.encoder.embed_tokens = self.shared
+ self.decoder.embed_tokens = self.shared
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: list[torch.FloatTensor] | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ decoder_inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Seq2SeqModelOutput:
+ r"""
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Indices of decoder input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+ BlenderbotSmall uses the `bos_token_id` as the starting token for `decoder_input_ids` generation. If
+ `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
+ `past_key_values`).
+ decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+ be used by default.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, BlenderbotSmallModel
+
+ >>> model = BlenderbotSmallModel.from_pretrained("facebook/blenderbot_small-90M")
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot_small-90M")
+
+ >>> inputs = tokenizer("Studies have been shown that owning a dog is good for you", return_tensors="pt")
+ >>> decoder_inputs = tokenizer("Studies show that", return_tensors="pt") # Batch size 1
+ >>> outputs = model(input_ids=inputs.input_ids, decoder_input_ids=decoder_inputs.input_ids)
+
+ >>> last_hidden_states = outputs.last_hidden_state
+ >>> list(last_hidden_states.shape)
+ [1, 3, 512]
+ ```"""
+
+ if encoder_outputs is None:
+ encoder_outputs: BaseModelOutput = self.encoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+ elif not isinstance(encoder_outputs, BaseModelOutput):
+ encoder_outputs = BaseModelOutput(
+ last_hidden_state=encoder_outputs[0],
+ hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,
+ attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,
+ )
+
+ decoder_outputs: BaseModelOutputWithPastAndCrossAttentions = self.decoder(
+ input_ids=decoder_input_ids,
+ attention_mask=decoder_attention_mask,
+ encoder_hidden_states=encoder_outputs[0],
+ encoder_attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=decoder_inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ return Seq2SeqModelOutput(
+ last_hidden_state=decoder_outputs.last_hidden_state,
+ past_key_values=decoder_outputs.past_key_values,
+ decoder_hidden_states=decoder_outputs.hidden_states,
+ decoder_attentions=decoder_outputs.attentions,
+ cross_attentions=decoder_outputs.cross_attentions,
+ encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+ encoder_hidden_states=encoder_outputs.hidden_states,
+ encoder_attentions=encoder_outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The BlenderbotSmall Model with a language modeling head. Can be used for summarization.
+ """
+)
+class BlenderbotSmallForConditionalGeneration(BlenderbotSmallPreTrainedModel, GenerationMixin):
+ base_model_prefix = "model"
+ _keys_to_ignore_on_load_missing = ["final_logits_bias"]
+ _tied_weights_keys = {
+ "lm_head.weight": "model.shared.weight",
+ }
+
+ def __init__(self, config: BlenderbotSmallConfig):
+ super().__init__(config)
+ self.model = BlenderbotSmallModel(config)
+ self.register_buffer("final_logits_bias", torch.zeros((1, self.model.shared.num_embeddings)))
+ self.lm_head = nn.Linear(config.d_model, self.model.shared.num_embeddings, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def resize_token_embeddings(
+ self, new_num_tokens: int, pad_to_multiple_of: int | None = None, mean_resizing: bool = True
+ ) -> nn.Embedding:
+ new_embeddings = super().resize_token_embeddings(new_num_tokens, pad_to_multiple_of, mean_resizing)
+ self._resize_final_logits_bias(new_embeddings.weight.shape[0])
+ return new_embeddings
+
+ def _resize_final_logits_bias(self, new_num_tokens: int) -> None:
+ old_num_tokens = self.final_logits_bias.shape[-1]
+ if new_num_tokens <= old_num_tokens:
+ new_bias = self.final_logits_bias[:, :new_num_tokens]
+ else:
+ extra_bias = torch.zeros((1, new_num_tokens - old_num_tokens), device=self.final_logits_bias.device)
+ new_bias = torch.cat([self.final_logits_bias, extra_bias], dim=1)
+ self.register_buffer("final_logits_bias", new_bias)
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ encoder_outputs: list[torch.FloatTensor] | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ decoder_inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> Seq2SeqLMOutput:
+ r"""
+ decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Indices of decoder input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are decoder input IDs?](../glossary#decoder-input-ids)
+
+ BlenderbotSmall uses the `bos_token_id` as the starting token for `decoder_input_ids` generation. If
+ `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
+ `past_key_values`).
+ decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+ be used by default.
+ 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]`.
+
+ Example Conversation:
+
+ ```python
+ >>> from transformers import AutoTokenizer, BlenderbotSmallForConditionalGeneration
+
+ >>> mname = "facebook/blenderbot_small-90M"
+ >>> model = BlenderbotSmallForConditionalGeneration.from_pretrained(mname)
+ >>> tokenizer = AutoTokenizer.from_pretrained(mname)
+ >>> UTTERANCE = "My friends are cool but they eat too many carbs."
+ >>> print("Human: ", UTTERANCE)
+ Human: My friends are cool but they eat too many carbs.
+
+ >>> inputs = tokenizer([UTTERANCE], return_tensors="pt")
+ >>> reply_ids = model.generate(**inputs)
+ >>> print("Bot: ", tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0])
+ Bot: what kind of carbs do they eat? i don't know much about carbs.
+
+ >>> REPLY = "I'm not sure"
+ >>> print("Human: ", REPLY)
+ Human: I'm not sure
+
+ >>> NEXT_UTTERANCE = (
+ ... "My friends are cool but they eat too many carbs.__end__ __start__what kind of carbs do they eat? "
+ ... "i don't know much about carbs__end__ "
+ ... "__start__ I'm not sure."
+ ... )
+ >>> inputs = tokenizer([NEXT_UTTERANCE], return_tensors="pt")
+ >>> next_reply_ids = model.generate(**inputs)
+ >>> print("Bot: ", tokenizer.batch_decode(next_reply_ids, skip_special_tokens=True)[0])
+ Bot: they eat a lot of carbs. carbs are high in fat, protein, and fats.
+ ```
+ """
+ if labels is not None:
+ if use_cache:
+ logger.warning("The `use_cache` argument is changed to `False` since `labels` is provided.")
+ use_cache = False
+ if decoder_input_ids is None and decoder_inputs_embeds is None:
+ decoder_input_ids = shift_tokens_right(
+ labels, self.config.pad_token_id, self.config.decoder_start_token_id
+ )
+
+ outputs: Seq2SeqModelOutput = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ decoder_input_ids=decoder_input_ids,
+ encoder_outputs=encoder_outputs,
+ decoder_attention_mask=decoder_attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ decoder_inputs_embeds=decoder_inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ lm_logits = self.lm_head(outputs[0])
+ lm_logits = lm_logits + self.final_logits_bias.to(lm_logits.device)
+
+ masked_lm_loss = None
+ if labels is not None:
+ labels = labels.to(lm_logits.device)
+ loss_fct = CrossEntropyLoss()
+ masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+ return Seq2SeqLMOutput(
+ loss=masked_lm_loss,
+ logits=lm_logits,
+ past_key_values=outputs.past_key_values,
+ decoder_hidden_states=outputs.decoder_hidden_states,
+ decoder_attentions=outputs.decoder_attentions,
+ cross_attentions=outputs.cross_attentions,
+ encoder_last_hidden_state=outputs.encoder_last_hidden_state,
+ encoder_hidden_states=outputs.encoder_hidden_states,
+ encoder_attentions=outputs.encoder_attentions,
+ )
+
+
+# Copied from transformers.models.bart.modeling_bart.BartDecoderWrapper with Bart->BlenderbotSmall
+class BlenderbotSmallDecoderWrapper(BlenderbotSmallPreTrainedModel):
+ """
+ This wrapper class is a helper class to correctly load pretrained checkpoints when the causal language model is
+ used in combination with the [`EncoderDecoderModel`] framework.
+ """
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.decoder = BlenderbotSmallDecoder(config)
+ self.post_init()
+
+ def forward(self, *args, **kwargs):
+ return self.decoder(*args, **kwargs)
+
+
+# Copied from transformers.models.bart.modeling_bart.BartForCausalLM with Bart->BlenderbotSmall, facebook/bart-base->facebook/blenderbot_small-90M
+class BlenderbotSmallForCausalLM(BlenderbotSmallPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {
+ "lm_head.weight": "model.decoder.embed_tokens.weight",
+ }
+
+ def __init__(self, config):
+ config.is_decoder = True
+ config.is_encoder_decoder = False
+ super().__init__(config)
+ self.model = BlenderbotSmallDecoderWrapper(config)
+
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.decoder.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.model.decoder.embed_tokens = value
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithCrossAttentions:
+ r"""
+ 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]`.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, BlenderbotSmallForCausalLM
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("facebook/blenderbot_small-90M")
+ >>> model = BlenderbotSmallForCausalLM.from_pretrained("facebook/blenderbot_small-90M")
+ >>> assert model.config.is_decoder, f"{model.__class__} has to be configured as a decoder."
+ >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+ >>> outputs = model(**inputs)
+
+ >>> logits = outputs.logits
+ >>> expected_shape = [1, inputs.input_ids.shape[-1], model.config.vocab_size]
+ >>> list(logits.shape) == expected_shape
+ True
+ ```"""
+
+ outputs: BaseModelOutputWithPastAndCrossAttentions = self.model.decoder(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ # Only compute necessary logits
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ labels = labels.to(logits.device)
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+ return CausalLMOutputWithCrossAttentions(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+__all__ = [
+ "BlenderbotSmallForCausalLM",
+ "BlenderbotSmallForConditionalGeneration",
+ "BlenderbotSmallModel",
+ "BlenderbotSmallPreTrainedModel",
+]
diff --git a/third_party/transformers/src/transformers/models/blenderbot_small/tokenization_blenderbot_small.py b/third_party/transformers/src/transformers/models/blenderbot_small/tokenization_blenderbot_small.py
new file mode 100644
index 0000000000000000000000000000000000000000..6d31cbe4432cfa9c75572e4265cba8a9a6830f13
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/blenderbot_small/tokenization_blenderbot_small.py
@@ -0,0 +1,222 @@
+# Copyright 2021 The Facebook Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Tokenization class for BlenderbotSmall."""
+
+import json
+import os
+
+import regex as re
+
+from ...tokenization_python import PreTrainedTokenizer
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+VOCAB_FILES_NAMES = {
+ "vocab_file": "vocab.json",
+ "merges_file": "merges.txt",
+ "tokenizer_config_file": "tokenizer_config.json",
+}
+
+
+def get_pairs(word):
+ """
+ Return set of symbol pairs in a word.
+
+ Word is represented as tuple of symbols (symbols being variable-length strings).
+ """
+ pairs = set()
+ prev_char = word[0]
+ for char in word[1:]:
+ pairs.add((prev_char, char))
+ prev_char = char
+
+ pairs = set(pairs)
+ return pairs
+
+
+class BlenderbotSmallTokenizer(PreTrainedTokenizer):
+ """
+ Constructs a Blenderbot-90M tokenizer based on BPE (Byte-Pair-Encoding)
+
+ This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+ the superclass for more information regarding methods.
+
+ Args:
+ vocab_file (`str`):
+ File containing the vocabulary.
+ merges_file (`str`):
+ Path to the merges file.
+ bos_token (`str`, *optional*, defaults to `"__start__"`):
+ The beginning of sentence token.
+ eos_token (`str`, *optional*, defaults to `"__end__"`):
+ The end of sentence token.
+ unk_token (`str`, *optional*, defaults to `"__unk__"`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ pad_token (`str`, *optional*, defaults to `"__null__"`):
+ The token used for padding, for example when batching sequences of different lengths.
+ kwargs (*optional*):
+ Additional keyword arguments passed along to [`PreTrainedTokenizer`]
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ model_input_names = ["input_ids", "attention_mask"]
+
+ def __init__(
+ self,
+ vocab_file,
+ merges_file,
+ bos_token="__start__",
+ eos_token="__end__",
+ unk_token="__unk__",
+ pad_token="__null__",
+ **kwargs,
+ ):
+ with open(vocab_file, encoding="utf-8") as vocab_handle:
+ self.encoder = json.load(vocab_handle)
+ self.decoder = {v: k for k, v in self.encoder.items()}
+ with open(merges_file, encoding="utf-8") as merges_handle:
+ merges = merges_handle.read().split("\n")[1:-1]
+ merges = [tuple(merge.split()) for merge in merges]
+ self.bpe_ranks = dict(zip(merges, range(len(merges))))
+ self.cache = {}
+
+ super().__init__(unk_token=unk_token, bos_token=bos_token, eos_token=eos_token, pad_token=pad_token, **kwargs)
+ self.special_tokens_pattern = None
+
+ @property
+ def vocab_size(self) -> int:
+ return len(self.encoder)
+
+ def get_vocab(self) -> dict:
+ return dict(self.encoder, **self.added_tokens_encoder)
+
+ def bpe(self, token: str) -> str:
+ if token in self.cache:
+ return self.cache[token]
+ token = re.sub("([.,!?()])", r" \1", token)
+ token = re.sub("(')", r" \1 ", token)
+ token = re.sub(r"\s{2,}", " ", token)
+ if "\n" in token:
+ token = token.replace("\n", " __newln__")
+
+ tokens = token.split(" ")
+ words = []
+ for token in tokens:
+ if not len(token):
+ continue
+
+ token = token.lower()
+ word = tuple(token)
+ word = tuple(list(word[:-1]) + [word[-1] + ""])
+ pairs = get_pairs(word)
+
+ if not pairs:
+ words.append(token)
+ continue
+
+ while True:
+ bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
+ if bigram not in self.bpe_ranks:
+ break
+ first, second = bigram
+ new_word = []
+ i = 0
+
+ while i < len(word):
+ try:
+ j = word.index(first, i)
+ new_word.extend(word[i:j])
+ i = j
+ except ValueError:
+ new_word.extend(word[i:])
+ break
+
+ if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
+ new_word.append(first + second)
+ i += 2
+ else:
+ new_word.append(word[i])
+ i += 1
+ new_word = tuple(new_word)
+ word = new_word
+ if len(word) == 1:
+ break
+ else:
+ pairs = get_pairs(word)
+ word = "@@ ".join(word)
+ word = word[:-4]
+
+ self.cache[token] = word
+ words.append(word)
+ return " ".join(words)
+
+ def _tokenize(self, text: str) -> list[str]:
+ """Split a string into tokens using BPE."""
+ split_tokens = []
+
+ words = re.findall(r"\S+\n?", text)
+
+ for token in words:
+ split_tokens.extend(list(self.bpe(token).split(" ")))
+ return split_tokens
+
+ def _convert_token_to_id(self, token: str) -> int:
+ """Converts a token to an id using the vocab."""
+ token = token.lower()
+ return self.encoder.get(token, self.encoder.get(self.unk_token))
+
+ def _convert_id_to_token(self, index: int) -> str:
+ """Converts an index (integer) in a token (str) using the vocab."""
+ return self.decoder.get(index, self.unk_token)
+
+ def convert_tokens_to_string(self, tokens: list[str]) -> str:
+ """Converts a sequence of tokens in a single string."""
+ out_string = " ".join(tokens).replace("@@ ", "").strip()
+ return out_string
+
+ def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
+ if not os.path.isdir(save_directory):
+ logger.error(f"Vocabulary path ({save_directory}) should be a directory")
+ return
+ vocab_file = os.path.join(
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+ )
+ merge_file = os.path.join(
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"]
+ )
+
+ with open(vocab_file, "w", encoding="utf-8") as f:
+ f.write(json.dumps(self.encoder, indent=2, sort_keys=True, ensure_ascii=False) + "\n")
+
+ index = 0
+ with open(merge_file, "w", encoding="utf-8") as writer:
+ writer.write("#version: 0.2\n")
+ for bpe_tokens, token_index in sorted(self.bpe_ranks.items(), key=lambda kv: kv[1]):
+ if index != token_index:
+ logger.warning(
+ f"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
+ " Please check that the tokenizer is not corrupted!"
+ )
+ index = token_index
+ writer.write(" ".join(bpe_tokens) + "\n")
+ index += 1
+
+ return vocab_file, merge_file
+
+
+__all__ = ["BlenderbotSmallTokenizer"]
diff --git a/third_party/transformers/src/transformers/models/bridgetower/__init__.py b/third_party/transformers/src/transformers/models/bridgetower/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..28a9840a14851e41c6966d840558b0fb34571605
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bridgetower/__init__.py
@@ -0,0 +1,30 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_bridgetower import *
+ from .image_processing_bridgetower import *
+ from .image_processing_pil_bridgetower import *
+ from .modeling_bridgetower import *
+ from .processing_bridgetower import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/bridgetower/configuration_bridgetower.py b/third_party/transformers/src/transformers/models/bridgetower/configuration_bridgetower.py
new file mode 100644
index 0000000000000000000000000000000000000000..d5be975c9aeadc3c65c422a195cea18a99bd41c5
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bridgetower/configuration_bridgetower.py
@@ -0,0 +1,167 @@
+# Copyright 2023 The Intel Labs Team Authors, The Microsoft Research Team Authors and HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""BridgeTower model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="BridgeTower/bridgetower-base")
+@strict
+class BridgeTowerVisionConfig(PreTrainedConfig):
+ r"""
+ stop_gradient (`bool`, *optional*, defaults to `False`):
+ Whether to stop gradient for training.
+ share_layernorm (`bool`, *optional*, defaults to `True`):
+ Whether LayerNorm layers are shared.
+ remove_last_layer (`bool`, *optional*, defaults to `False`):
+ Whether to remove the last layer from the vision encoder.
+
+ Example:
+
+ ```python
+ >>> from transformers import BridgeTowerVisionConfig
+
+ >>> # Initializing a BridgeTower BridgeTower/bridgetower-base style configuration for the vision model
+ >>> configuration = BridgeTowerVisionConfig()
+
+ >>> # Accessing the configuration
+ >>> configuration
+ ```"""
+
+ model_type = "bridgetower_vision_model"
+ base_config_key = "vision_config"
+
+ hidden_size: int = 768
+ num_hidden_layers: int = 12
+ num_channels: int = 3
+ patch_size: int | list[int] | tuple[int, int] = 16
+ image_size: int | list[int] | tuple[int, int] = 288
+ initializer_factor: float | int = 1
+ layer_norm_eps: float = 1e-05
+ stop_gradient: bool = False
+ share_layernorm: bool = True
+ remove_last_layer: bool = False
+
+
+@auto_docstring(checkpoint="BridgeTower/bridgetower-base")
+@strict
+class BridgeTowerTextConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import BridgeTowerTextConfig
+
+ >>> # Initializing a BridgeTower BridgeTower/bridgetower-base style configuration for the text model
+ >>> configuration = BridgeTowerTextConfig()
+
+ >>> # Accessing the configuration
+ >>> configuration
+ ```"""
+
+ model_type = "bridgetower_text_model"
+ base_config_key = "text_config"
+
+ vocab_size: int = 50265
+ hidden_size: int = 768
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 12
+ initializer_factor: float | int = 1
+ intermediate_size: int = 3072
+ hidden_act: str = "gelu"
+ hidden_dropout_prob: float | int = 0.1
+ attention_probs_dropout_prob: float | int = 0.1
+ max_position_embeddings: int = 514
+ type_vocab_size: int = 1
+ layer_norm_eps: float = 1e-05
+ pad_token_id: int | None = 1
+ bos_token_id: int | None = 0
+ eos_token_id: int | list[int] | None = 2
+ use_cache: bool = True
+ is_decoder: bool = False
+ add_cross_attention: bool = False
+
+
+@auto_docstring(checkpoint="BridgeTower/bridgetower-base")
+@strict
+class BridgeTowerConfig(PreTrainedConfig):
+ r"""
+ share_cross_modal_transformer_layers (`bool`, *optional*, defaults to `True`):
+ Whether cross modal transformer layers are shared.
+ share_link_tower_layers (`bool`, *optional*, defaults to `False`):
+ Whether the bride/link tower layers are shared.
+ link_tower_type (`str`, *optional*, defaults to `"add"`):
+ Type of the bridge/link layer.
+ init_layernorm_from_vision_encoder (`bool`, *optional*, defaults to `False`):
+ Whether to init LayerNorm from the vision encoder.
+
+ Example:
+
+ ```python
+ >>> from transformers import BridgeTowerModel, BridgeTowerConfig
+
+ >>> # Initializing a BridgeTower BridgeTower/bridgetower-base style configuration
+ >>> configuration = BridgeTowerConfig()
+
+ >>> # Initializing a model from the BridgeTower/bridgetower-base style configuration
+ >>> model = BridgeTowerModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "bridgetower"
+ sub_configs = {"text_config": BridgeTowerTextConfig, "vision_config": BridgeTowerVisionConfig}
+
+ share_cross_modal_transformer_layers: bool = True
+ hidden_act: str = "gelu"
+ hidden_size: int = 768
+ initializer_factor: float | int = 1
+ layer_norm_eps: float = 1e-05
+ share_link_tower_layers: bool = False
+ link_tower_type: str = "add"
+ num_attention_heads: int = 12
+ num_hidden_layers: int = 6
+ tie_word_embeddings: bool = False
+ init_layernorm_from_vision_encoder: bool = False
+ text_config: dict | PreTrainedConfig | None = None
+ vision_config: dict | PreTrainedConfig | None = None
+
+ def __post_init__(self, **kwargs):
+ # TODO: remove this once the Hub files are updated.
+ _ = kwargs.pop("text_config_dict", None)
+ _ = kwargs.pop("vision_config_dict", None)
+
+ if self.text_config is None:
+ self.text_config = BridgeTowerTextConfig()
+ logger.info("`text_config` is `None`. initializing the `BridgeTowerTextConfig` with default values.")
+ elif isinstance(self.text_config, dict):
+ self.text_config = BridgeTowerTextConfig(**self.text_config)
+
+ if self.vision_config is None:
+ self.vision_config = BridgeTowerVisionConfig()
+ logger.info("`vision_config` is `None`. initializing the `BridgeTowerVisionConfig` with default values.")
+ elif isinstance(self.vision_config, dict):
+ self.vision_config = BridgeTowerVisionConfig(**self.vision_config)
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["BridgeTowerConfig", "BridgeTowerTextConfig", "BridgeTowerVisionConfig"]
diff --git a/third_party/transformers/src/transformers/models/bridgetower/image_processing_bridgetower.py b/third_party/transformers/src/transformers/models/bridgetower/image_processing_bridgetower.py
new file mode 100644
index 0000000000000000000000000000000000000000..1a20dd8316751ae597c1bb83f79940f55fea40e1
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bridgetower/image_processing_bridgetower.py
@@ -0,0 +1,184 @@
+# Copyright 2023 The Intel Labs Team Authors, The Microsoft Research Team Authors and HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Image processor class for BridgeTower."""
+
+from typing import Union
+
+import numpy as np
+import torch
+from torchvision.transforms.v2 import functional as tvF
+
+from ...image_processing_backends import TorchvisionBackend
+from ...image_processing_utils import BatchFeature
+from ...image_transforms import group_images_by_shape, reorder_images
+from ...image_utils import (
+ OPENAI_CLIP_MEAN,
+ OPENAI_CLIP_STD,
+ ChannelDimension,
+ PILImageResampling,
+ SizeDict,
+ get_image_size,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring, is_torch_available
+
+
+def get_resize_output_image_size(
+ input_image: Union[np.ndarray, "torch.Tensor"],
+ shorter: int = 800,
+ longer: int = 1333,
+ size_divisor: int = 32,
+) -> tuple[int, int]:
+ """Get output image size after resizing with size_divisor."""
+ if is_torch_available() and isinstance(input_image, torch.Tensor):
+ input_height, input_width = input_image.shape[-2:]
+ else:
+ input_height, input_width = get_image_size(input_image, channel_dim=ChannelDimension.FIRST)
+
+ min_size, max_size = shorter, longer
+ scale = min_size / min(input_height, input_width)
+
+ if input_height < input_width:
+ new_height = min_size
+ new_width = scale * input_width
+ else:
+ new_height = scale * input_height
+ new_width = min_size
+
+ if max(new_height, new_width) > max_size:
+ scale = max_size / max(new_height, new_width)
+ new_height = scale * new_height
+ new_width = scale * new_width
+
+ new_height, new_width = int(new_height + 0.5), int(new_width + 0.5)
+ new_height = new_height // size_divisor * size_divisor
+ new_width = new_width // size_divisor * size_divisor
+
+ return new_height, new_width
+
+
+class BridgeTowerImageProcessorKwargs(ImagesKwargs, total=False):
+ r"""
+ size_divisor (`int`, *optional*, defaults to `self.size_divisor`):
+ The size by which to make sure both the height and width can be divided.
+ """
+
+ size_divisor: int
+
+
+@auto_docstring
+class BridgeTowerImageProcessor(TorchvisionBackend):
+ """Torchvision backend for BridgeTower with custom resize and center_crop."""
+
+ valid_kwargs = BridgeTowerImageProcessorKwargs
+ model_input_names = ["pixel_values", "pixel_mask"]
+
+ resample = PILImageResampling.BICUBIC
+ image_mean = OPENAI_CLIP_MEAN
+ image_std = OPENAI_CLIP_STD
+ size = {"shortest_edge": 288}
+ default_to_square = False
+ crop_size = {"shortest_edge": 288}
+ do_resize = True
+ do_center_crop = True
+ do_rescale = True
+ do_normalize = True
+ do_pad = True
+ size_divisor = 32
+
+ def __init__(self, **kwargs: Unpack[BridgeTowerImageProcessorKwargs]):
+ super().__init__(**kwargs)
+
+ def resize(
+ self,
+ image: "torch.Tensor",
+ size: SizeDict,
+ resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+ size_divisor: int = 32,
+ **kwargs,
+ ) -> "torch.Tensor":
+ """Resize with size_divisor support."""
+ if not size.shortest_edge:
+ raise ValueError(f"The `size` dictionary must contain the key `shortest_edge`. Got {size.keys()}")
+ shorter = size.shortest_edge
+ longer = int(1333 / 800 * shorter)
+ output_height, output_width = get_resize_output_image_size(
+ image, shorter=shorter, longer=longer, size_divisor=size_divisor
+ )
+ return super().resize(
+ image=image,
+ size=SizeDict(height=output_height, width=output_width),
+ resample=resample,
+ **kwargs,
+ )
+
+ def _preprocess(
+ self,
+ images: list["torch.Tensor"],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+ do_center_crop: bool,
+ crop_size: SizeDict,
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ do_pad: bool | None,
+ pad_size: SizeDict | None,
+ disable_grouping: bool | None,
+ return_tensors: str | TensorType | None,
+ size_divisor: int = 32,
+ **kwargs,
+ ) -> BatchFeature:
+ """Custom preprocessing for BridgeTower."""
+ # Group images by size for batched resizing
+ grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+ resized_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ if do_resize:
+ stacked_images = self.resize(stacked_images, size=size, resample=resample, size_divisor=size_divisor)
+ resized_images_grouped[shape] = stacked_images
+ resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+ # Group images by size for further processing
+ grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
+ processed_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ if do_center_crop:
+ stacked_images = self.center_crop(
+ stacked_images, size=SizeDict(height=crop_size.shortest_edge, width=crop_size.shortest_edge)
+ )
+ # Use fused rescale and normalize
+ stacked_images = self.rescale_and_normalize(
+ stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
+ )
+ processed_images_grouped[shape] = stacked_images
+
+ processed_images = reorder_images(processed_images_grouped, grouped_images_index)
+
+ data = {}
+ if do_pad:
+ processed_images, processed_masks = self.pad(
+ processed_images, return_mask=True, disable_grouping=disable_grouping
+ )
+ data["pixel_mask"] = processed_masks
+
+ data["pixel_values"] = processed_images
+
+ return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["BridgeTowerImageProcessor"]
diff --git a/third_party/transformers/src/transformers/models/bridgetower/image_processing_pil_bridgetower.py b/third_party/transformers/src/transformers/models/bridgetower/image_processing_pil_bridgetower.py
new file mode 100644
index 0000000000000000000000000000000000000000..e6d92a4b3f1373979d3a2a5a6358dfffd650c4ac
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bridgetower/image_processing_pil_bridgetower.py
@@ -0,0 +1,168 @@
+# Copyright 2023 The Intel Labs Team Authors, The Microsoft Research Team Authors and HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Image processor class for BridgeTower."""
+
+import numpy as np
+
+from ...image_processing_backends import PilBackend
+from ...image_processing_utils import BatchFeature
+from ...image_utils import (
+ OPENAI_CLIP_MEAN,
+ OPENAI_CLIP_STD,
+ ChannelDimension,
+ PILImageResampling,
+ SizeDict,
+ get_image_size,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring
+
+
+# Adapted from transformers.models.bridgetower.image_processing_bridgetower.BridgeTowerImageProcessorKwargs
+class BridgeTowerImageProcessorKwargs(ImagesKwargs, total=False):
+ r"""
+ size_divisor (`int`, *optional*, defaults to `self.size_divisor`):
+ The size by which to make sure both the height and width can be divided.
+ """
+
+ size_divisor: int
+
+
+# adapted from transformers.models.bridgetower.image_processing_bridgetower.get_resize_output_image_size
+def get_resize_output_image_size(
+ input_image: np.ndarray,
+ shorter: int = 800,
+ longer: int = 1333,
+ size_divisor: int = 32,
+) -> tuple[int, int]:
+ """Get output image size after resizing with size_divisor."""
+ input_height, input_width = get_image_size(input_image, channel_dim=ChannelDimension.FIRST)
+
+ min_size, max_size = shorter, longer
+ scale = min_size / min(input_height, input_width)
+
+ if input_height < input_width:
+ new_height = min_size
+ new_width = scale * input_width
+ else:
+ new_height = scale * input_height
+ new_width = min_size
+
+ if max(new_height, new_width) > max_size:
+ scale = max_size / max(new_height, new_width)
+ new_height = scale * new_height
+ new_width = scale * new_width
+
+ new_height, new_width = int(new_height + 0.5), int(new_width + 0.5)
+ new_height = new_height // size_divisor * size_divisor
+ new_width = new_width // size_divisor * size_divisor
+
+ return new_height, new_width
+
+
+@auto_docstring
+class BridgeTowerImageProcessorPil(PilBackend):
+ """PIL backend for BridgeTower with custom resize and center_crop."""
+
+ valid_kwargs = BridgeTowerImageProcessorKwargs
+ model_input_names = ["pixel_values", "pixel_mask"]
+
+ resample = PILImageResampling.BICUBIC
+ image_mean = OPENAI_CLIP_MEAN
+ image_std = OPENAI_CLIP_STD
+ size = {"shortest_edge": 288}
+ default_to_square = False
+ crop_size = {"shortest_edge": 288}
+ do_resize = True
+ do_center_crop = True
+ do_rescale = True
+ do_normalize = True
+ do_pad = True
+ size_divisor = 32
+
+ def __init__(self, **kwargs: Unpack[BridgeTowerImageProcessorKwargs]):
+ super().__init__(**kwargs)
+
+ def resize(
+ self,
+ image: np.ndarray,
+ size: SizeDict,
+ resample: "PILImageResampling | None",
+ size_divisor: int = 32,
+ **kwargs,
+ ) -> np.ndarray:
+ """Resize with size_divisor support."""
+
+ if not size.shortest_edge:
+ raise ValueError(f"The `size` dictionary must contain the key `shortest_edge`. Got {size.keys()}")
+ shorter = size.shortest_edge
+ longer = int(1333 / 800 * shorter)
+ output_height, output_width = get_resize_output_image_size(
+ image, shorter=shorter, longer=longer, size_divisor=size_divisor
+ )
+ return super().resize(
+ image,
+ size=SizeDict(height=output_height, width=output_width),
+ resample=resample,
+ **kwargs,
+ )
+
+ def _preprocess(
+ self,
+ images: list[np.ndarray],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | None",
+ do_center_crop: bool,
+ crop_size: SizeDict,
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ do_pad: bool | None,
+ pad_size: SizeDict | None,
+ return_tensors: str | TensorType | None,
+ size_divisor: int = 32,
+ **kwargs,
+ ) -> BatchFeature:
+ """Custom preprocessing for BridgeTower."""
+ processed_images = []
+ for image in images:
+ if do_resize:
+ image = self.resize(image, size, resample, size_divisor)
+ if do_center_crop:
+ image = self.center_crop(
+ image, size=SizeDict(height=crop_size.shortest_edge, width=crop_size.shortest_edge)
+ )
+ if do_rescale:
+ image = self.rescale(image, rescale_factor)
+ if do_normalize:
+ image = self.normalize(image, image_mean, image_std)
+ processed_images.append(image)
+
+ data = {}
+ if do_pad:
+ processed_images, processed_masks = self.pad(
+ processed_images,
+ return_mask=True,
+ )
+ data["pixel_mask"] = processed_masks
+
+ data["pixel_values"] = processed_images
+
+ return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["BridgeTowerImageProcessorPil"]
diff --git a/third_party/transformers/src/transformers/models/bridgetower/modeling_bridgetower.py b/third_party/transformers/src/transformers/models/bridgetower/modeling_bridgetower.py
new file mode 100644
index 0000000000000000000000000000000000000000..81bbf8797129e1f861bb8fbc5b4b93369591b384
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bridgetower/modeling_bridgetower.py
@@ -0,0 +1,1762 @@
+# Copyright 2023 The Intel Labs Team Authors, The Microsoft Research Team Authors and HuggingFace Inc. team. All rights reserved.
+#
+# 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 BridgeTower Model"""
+
+from collections import OrderedDict
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from ... import initialization as init
+from ...activations import ACT2FN, QuickGELUActivation
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...masking_utils import create_bidirectional_mask, create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutputWithPastAndCrossAttentions,
+ BaseModelOutputWithPoolingAndCrossAttentions,
+ MaskedLMOutput,
+ ModelOutput,
+ SequenceClassifierOutput,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...pytorch_utils import apply_chunking_to_forward
+from ...utils import TransformersKwargs, auto_docstring, logging, torch_int
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_bridgetower import BridgeTowerConfig, BridgeTowerTextConfig, BridgeTowerVisionConfig
+
+
+logger = logging.get_logger(__name__)
+
+_TOKENIZER_FOR_DOC = "RobertaTokenizer"
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Output type of [`BridgeTowerModel`].
+ """
+)
+class BridgeTowerModelOutput(ModelOutput):
+ r"""
+ text_features (`torch.FloatTensor` of shape `(batch_size, text_sequence_length, hidden_size)`):
+ Sequence of hidden-states at the text output of the last layer of the model.
+ image_features (`torch.FloatTensor` of shape `(batch_size, image_sequence_length, hidden_size)`):
+ Sequence of hidden-states at the image output of the last layer of the model.
+ pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size x 2)`):
+ Concatenation of last layer hidden-state of the first token of the text and image sequence (classification
+ token), respectively, after further processing through layers used for auxiliary pretraining tasks.
+ """
+
+ text_features: torch.FloatTensor | None = None
+ image_features: torch.FloatTensor | None = None
+ pooler_output: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Output type of ['BridgeTowerForContrastiveLearning']
+ """
+)
+class BridgeTowerContrastiveOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`):
+ Image-text contrastive loss.
+ 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).
+ text_embeds (`torch.FloatTensor)`, *optional*, returned when model is initialized with `with_projection=True`):
+ The text embeddings obtained by applying the projection layer to the pooler_output.
+ image_embeds (`torch.FloatTensor)`, *optional*, returned when model is initialized with `with_projection=True`):
+ The image embeddings obtained by applying the projection layer to the pooler_output.
+ cross_embeds (`torch.FloatTensor)`, *optional*, returned when model is initialized with `with_projection=True`):
+ The text-image cross-modal embeddings obtained by applying the projection layer to the pooler_output.
+ 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)`.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ text_embeds: tuple[torch.FloatTensor] | None = None
+ image_embeds: tuple[torch.FloatTensor] | None = None
+ cross_embeds: tuple[torch.FloatTensor] | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+
+
+class BridgeTowerResidualAttention(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ self.attn = nn.MultiheadAttention(config.hidden_size, config.hidden_size // 64)
+ self.ln_1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.mlp = nn.ModuleDict(
+ OrderedDict(
+ [
+ ("c_fc", nn.Linear(config.hidden_size, config.hidden_size * 4)),
+ ("gelu", QuickGELUActivation()),
+ ("c_proj", nn.Linear(config.hidden_size * 4, config.hidden_size)),
+ ]
+ )
+ )
+ self.ln_2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.attn_mask = None
+
+ def attention(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor):
+ if attention_mask is not None:
+ attention_mask = attention_mask.to(dtype=torch.bool, device=hidden_state.device)
+ self.attn_mask = (
+ self.attn_mask.to(dtype=hidden_state.dtype, device=hidden_state.device)
+ if self.attn_mask is not None
+ else None
+ )
+ return self.attn(
+ hidden_state,
+ hidden_state,
+ hidden_state,
+ need_weights=False,
+ attn_mask=self.attn_mask,
+ key_padding_mask=attention_mask,
+ )[0]
+
+ def forward(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor | None = None):
+ residual_state = hidden_state + self.attention(self.ln_1(hidden_state), attention_mask)
+ hidden_state = self.ln_2(residual_state)
+ for layer in self.mlp.values():
+ hidden_state = layer(hidden_state)
+ hidden_state = residual_state + hidden_state
+ return hidden_state
+
+
+class BridgeTowerTransformer(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.num_hidden_layers = config.num_hidden_layers
+ if config.remove_last_layer:
+ self.resblocks = nn.ModuleList(
+ [BridgeTowerResidualAttention(config) for _ in range(self.num_hidden_layers - 1)]
+ )
+ else:
+ self.resblocks = nn.ModuleList(
+ [BridgeTowerResidualAttention(config) for _ in range(self.num_hidden_layers)]
+ )
+ self.stop_gradient = config.stop_gradient
+
+ def forward(self, hidden_state: torch.Tensor, attention_mask: torch.Tensor | None = None):
+ hidden_states = []
+ for block in self.resblocks:
+ hidden_state = block(hidden_state, attention_mask)
+ if self.stop_gradient:
+ hidden_states.append(hidden_state.detach())
+ else:
+ hidden_states.append(hidden_state)
+ return hidden_states
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPVisionEmbeddings with CLIP->BridgeTower
+class BridgeTowerVisionEmbeddings(nn.Module):
+ def __init__(self, config: BridgeTowerVisionConfig):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.image_size = config.image_size
+ self.patch_size = config.patch_size
+
+ self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))
+
+ self.patch_embedding = nn.Conv2d(
+ in_channels=config.num_channels,
+ out_channels=self.embed_dim,
+ kernel_size=self.patch_size,
+ stride=self.patch_size,
+ bias=False,
+ )
+
+ self.num_patches = (self.image_size // self.patch_size) ** 2
+ self.num_positions = self.num_patches + 1
+ self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+ self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)
+
+ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+ """
+ This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution
+ images. This method is also adapted to support torch.jit tracing.
+
+ Adapted from:
+ - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
+ - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
+ """
+
+ num_patches = embeddings.shape[1] - 1
+ position_embedding = self.position_embedding.weight.unsqueeze(0)
+ num_positions = position_embedding.shape[1] - 1
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embedding(self.position_ids)
+
+ class_pos_embed = position_embedding[:, :1]
+ patch_pos_embed = position_embedding[:, 1:]
+
+ dim = embeddings.shape[-1]
+
+ new_height = height // self.patch_size
+ new_width = width // self.patch_size
+
+ sqrt_num_positions = torch_int(num_positions**0.5)
+ patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
+ patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed,
+ size=(new_height, new_width),
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+ return torch.cat((class_pos_embed, patch_pos_embed), dim=1)
+
+ def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=False) -> torch.Tensor:
+ batch_size, _, height, width = pixel_values.shape
+ if not interpolate_pos_encoding and (height != self.image_size or width != self.image_size):
+ raise ValueError(
+ f"Input image size ({height}*{width}) doesn't match model ({self.image_size}*{self.image_size})."
+ )
+ target_dtype = self.patch_embedding.weight.dtype
+ patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid]
+ patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+
+ class_embeds = self.class_embedding.expand(batch_size, 1, -1)
+ embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
+ if interpolate_pos_encoding:
+ embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+ else:
+ embeddings = embeddings + self.position_embedding(self.position_ids)
+ return embeddings
+
+
+class BridgeTowerVisionTransformer(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ self.embeddings = BridgeTowerVisionEmbeddings(config)
+ self.ln_pre = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.transformer = BridgeTowerTransformer(config)
+ self.ln_post = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.share_layernorm = config.share_layernorm
+ if not config.share_layernorm:
+ self.ln_separate = nn.ModuleList(
+ [nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) for _ in range(config.num_hidden_layers)]
+ )
+
+ def forward(
+ self,
+ pixel_values: torch.Tensor,
+ attention_mask,
+ interpolate_pos_encoding: bool = False,
+ ):
+ hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding)
+ hidden_states = self.ln_pre(hidden_states)
+ # NLD -> LND
+ hidden_states = hidden_states.permute(1, 0, 2)
+
+ hidden_states = self.transformer(hidden_states, attention_mask)
+ # shape = [num_hidden_layers, hidden_size, *, grid ** 2]
+ hidden_states = torch.stack(hidden_states, dim=0)
+ # shape = [num_hidden_layers, *, hidden_size, grid ** 2]
+ hidden_states = hidden_states.permute(0, 2, 1, 3)
+ if self.share_layernorm:
+ hidden_states = self.ln_post(hidden_states)
+ else:
+ hidden_states_stack = []
+ for hidden_states, ln in zip(hidden_states, self.ln_separate):
+ hidden_states = ln(hidden_states)
+ hidden_states_stack.append(hidden_states)
+ # shape = [num_hidden_layers, *, hidden_size, grid ** 2]
+ hidden_states = torch.stack(hidden_states_stack, dim=0)
+ return hidden_states
+
+ def forward_pre(
+ self,
+ pixel_values: torch.Tensor,
+ interpolate_pos_encoding: bool = False,
+ ):
+ hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
+ hidden_states = self.ln_pre(hidden_states)
+ # NLD -> LND
+ hidden_states = hidden_states.permute(1, 0, 2)
+ return hidden_states
+
+ def forward_post(self, hidden_state: torch.Tensor):
+ visual_output_post = hidden_state.permute(1, 0, 2)
+ visual_output_post = self.ln_post(visual_output_post)
+ return visual_output_post
+
+
+class BridgeTowerLinkTower(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.link_tower_type = config.link_tower_type
+ self.hidden_size = config.hidden_size
+ if config.link_tower_type in ["add", "scaled_add", "interpolate"]:
+ if config.link_tower_type == "scaled_add":
+ self.scaled_factor = nn.Parameter(torch.tensor(1.0))
+ elif config.link_tower_type == "interpolate":
+ self.beta = nn.Parameter(torch.tensor(0.5))
+ self.LayerNorm = nn.LayerNorm(self.hidden_size, eps=config.layer_norm_eps)
+ else:
+ raise NotImplementedError(f"link_tower_type {config.link_tower_type} is not implemented")
+
+ def forward(self, hidden_states, cross_modal_hidden_states, attention_mask):
+ if self.link_tower_type == "add":
+ return self.LayerNorm(hidden_states + cross_modal_hidden_states)
+ elif self.link_tower_type == "scaled_add":
+ return self.LayerNorm(hidden_states * self.scaled_factor + cross_modal_hidden_states)
+ elif self.link_tower_type == "interpolate":
+ return self.LayerNorm(hidden_states * (1 - self.beta) + cross_modal_hidden_states * self.beta)
+ else:
+ raise NotImplementedError(f"link_tower_type {self.link_tower_type} is not implemented")
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfOutput with Bert->BridgeTower
+class BridgeTowerSelfOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->BridgeTower
+class BridgeTowerIntermediate(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->BridgeTower
+class BridgeTowerOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->BridgeTower
+class BridgeTowerPooler(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.activation = nn.Tanh()
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ # We "pool" the model by simply taking the hidden state corresponding
+ # to the first token.
+ first_token_tensor = hidden_states[:, 0]
+ pooled_output = self.dense(first_token_tensor)
+ pooled_output = self.activation(pooled_output)
+ return pooled_output
+
+
+# Copied from transformers.models.bert.modeling_bert.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float | None = None,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ if scaling is None:
+ scaling = query.size(-1) ** -0.5
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.roberta.modeling_roberta.RobertaSelfAttention with Roberta->BridgeTower
+class BridgeTowerSelfAttention(nn.Module):
+ def __init__(self, config, is_causal=False, layer_idx=None):
+ super().__init__()
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+ raise ValueError(
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+ f"heads ({config.num_attention_heads})"
+ )
+ self.config = config
+
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+ self.scaling = self.attention_head_size**-0.5
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ self.is_decoder = config.is_decoder
+ self.is_causal = is_causal
+ self.layer_idx = layer_idx
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.attention_head_size)
+
+ # get all proj
+ query_layer = self.query(hidden_states).view(*hidden_shape).transpose(1, 2)
+ key_layer = self.key(hidden_states).view(*hidden_shape).transpose(1, 2)
+ value_layer = self.value(hidden_states).view(*hidden_shape).transpose(1, 2)
+
+ if past_key_values is not None:
+ # decoder-only roberta can have a simple dynamic cache for example
+ current_past_key_values = past_key_values
+ if isinstance(past_key_values, EncoderDecoderCache):
+ current_past_key_values = past_key_values.self_attention_cache
+
+ # save all key/value_layer to cache to be re-used for fast auto-regressive generation
+ key_layer, value_layer = current_past_key_values.update(key_layer, value_layer, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_layer,
+ key_layer,
+ value_layer,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout.p,
+ scaling=self.scaling,
+ **kwargs,
+ )
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.roberta.modeling_roberta.RobertaCrossAttention with Roberta->BridgeTower
+class BridgeTowerCrossAttention(nn.Module):
+ def __init__(self, config, is_causal=False, layer_idx=None):
+ super().__init__()
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+ raise ValueError(
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+ f"heads ({config.num_attention_heads})"
+ )
+ self.config = config
+
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+ self.scaling = self.attention_head_size**-0.5
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ self.is_causal = is_causal
+ self.layer_idx = layer_idx
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ past_key_values: EncoderDecoderCache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor]:
+ # determine input shapes
+ bsz, tgt_len = hidden_states.shape[:-1]
+ src_len = encoder_hidden_states.shape[1]
+
+ q_input_shape = (bsz, tgt_len, -1, self.attention_head_size)
+ kv_input_shape = (bsz, src_len, -1, self.attention_head_size)
+
+ # get query proj
+ query_layer = self.query(hidden_states).view(*q_input_shape).transpose(1, 2)
+
+ is_updated = past_key_values.is_updated.get(self.layer_idx) if past_key_values is not None else False
+ if past_key_values is not None and is_updated:
+ # reuse k,v, cross_attentions
+ key_layer = past_key_values.cross_attention_cache.layers[self.layer_idx].keys
+ value_layer = past_key_values.cross_attention_cache.layers[self.layer_idx].values
+ else:
+ key_layer = self.key(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2)
+ value_layer = self.value(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2)
+
+ if past_key_values is not None:
+ # save all states to the cache
+ key_layer, value_layer = past_key_values.cross_attention_cache.update(
+ key_layer, value_layer, self.layer_idx
+ )
+ # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+ past_key_values.is_updated[self.layer_idx] = True
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_layer,
+ key_layer,
+ value_layer,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout.p,
+ scaling=self.scaling,
+ **kwargs,
+ )
+ attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->BridgeTower,BERT->BRIDGE_TOWER
+class BridgeTowerAttention(nn.Module):
+ def __init__(self, config, is_causal=False, layer_idx=None, is_cross_attention=False):
+ super().__init__()
+ self.is_cross_attention = is_cross_attention
+ attention_class = BridgeTowerCrossAttention if is_cross_attention else BridgeTowerSelfAttention
+ self.self = attention_class(config, is_causal=is_causal, layer_idx=layer_idx)
+ self.output = BridgeTowerSelfOutput(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor]:
+ attention_mask = attention_mask if not self.is_cross_attention else encoder_attention_mask
+ attention_output, attn_weights = self.self(
+ hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ attention_output = self.output(attention_output, hidden_states)
+ return attention_output, attn_weights
+
+
+class BridgeTowerBertCrossLayer(nn.Module):
+ def __init__(self, config, layer_idx=None):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+ self.attention = BridgeTowerAttention(config, is_causal=True, layer_idx=layer_idx)
+ self.is_decoder = config.is_decoder
+ self.add_cross_attention = config.add_cross_attention
+ self.crossattention = BridgeTowerAttention(
+ config,
+ is_causal=False,
+ layer_idx=layer_idx,
+ is_cross_attention=True,
+ )
+ self.intermediate = BridgeTowerIntermediate(config)
+ self.output = BridgeTowerOutput(config)
+
+ def forward(
+ self,
+ hidden_states,
+ encoder_hidden_states,
+ attention_mask=None,
+ encoder_attention_mask=None,
+ past_key_values=None,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ self_attention_output, self_attn_weights = self.attention(
+ hidden_states,
+ attention_mask=attention_mask,
+ past_key_values=None,
+ **kwargs,
+ )
+ attention_output = self_attention_output
+
+ cross_attention_output, cross_attn_weights = self.crossattention(
+ attention_output,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ attention_output = cross_attention_output
+
+ layer_output = apply_chunking_to_forward(
+ self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+ )
+ return (
+ layer_output,
+ self_attn_weights,
+ cross_attn_weights,
+ )
+
+ def feed_forward_chunk(self, attention_output):
+ intermediate_output = self.intermediate(attention_output)
+ layer_output = self.output(intermediate_output, attention_output)
+ return layer_output
+
+
+class BridgeTowerTextLayer(GradientCheckpointingLayer):
+ def __init__(self, config, layer_idx=None):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+ self.attention = BridgeTowerAttention(config, is_causal=config.is_decoder, layer_idx=layer_idx)
+ self.is_decoder = config.is_decoder
+ self.add_cross_attention = config.add_cross_attention
+ if self.add_cross_attention:
+ if not self.is_decoder:
+ raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
+ self.crossattention = BridgeTowerAttention(
+ config,
+ is_causal=False,
+ layer_idx=layer_idx,
+ is_cross_attention=True,
+ )
+ self.intermediate = BridgeTowerIntermediate(config)
+ self.output = BridgeTowerOutput(config)
+
+ # copied from transformers.models.bert.modeling_bert.BertLayer.forward
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ self_attention_output, _ = self.attention(
+ hidden_states,
+ attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ attention_output = self_attention_output
+
+ if self.is_decoder and encoder_hidden_states is not None:
+ if not hasattr(self, "crossattention"):
+ raise ValueError(
+ f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
+ " by setting `config.add_cross_attention=True`"
+ )
+
+ cross_attention_output, _ = self.crossattention(
+ self_attention_output,
+ None, # attention_mask
+ encoder_hidden_states,
+ encoder_attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ attention_output = cross_attention_output
+
+ layer_output = apply_chunking_to_forward(
+ self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+ )
+ return layer_output
+
+ def feed_forward_chunk(self, attention_output):
+ intermediate_output = self.intermediate(attention_output)
+ layer_output = self.output(intermediate_output, attention_output)
+ return layer_output
+
+
+# copied from transformers.models.roberta.modeling_roberta.RobertaEncoder with Roberta->BridgeTowerText
+class BridgeTowerTextEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList(
+ [BridgeTowerTextLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)]
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPastAndCrossAttentions:
+ for layer_module in self.layer:
+ hidden_states = layer_module(
+ hidden_states,
+ attention_mask,
+ encoder_hidden_states, # as a positional argument for gradient checkpointing
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values if use_cache else None,
+ )
+
+
+# Copied from transformers.models.roberta.modeling_roberta.RobertaEmbeddings with Roberta->BridgeTowerText
+class BridgeTowerTextEmbeddings(nn.Module):
+ """Construct the embeddings from word, position and token_type embeddings."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+ self.register_buffer(
+ "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+ )
+ self.register_buffer(
+ "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
+ )
+
+ self.padding_idx = config.pad_token_id
+ self.position_embeddings = nn.Embedding(
+ config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
+ )
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values_length: int = 0,
+ ) -> torch.Tensor:
+ if position_ids is None:
+ if input_ids is not None:
+ # Create the position ids from the input token ids. Any padded tokens remain padded.
+ position_ids = self.create_position_ids_from_input_ids(
+ input_ids, self.padding_idx, past_key_values_length
+ )
+ else:
+ position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds, self.padding_idx)
+
+ if input_ids is not None:
+ input_shape = input_ids.size()
+ else:
+ input_shape = inputs_embeds.size()[:-1]
+
+ batch_size, seq_length = input_shape
+
+ # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
+ # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
+ # issue #5664
+ if token_type_ids is None:
+ if hasattr(self, "token_type_ids"):
+ # NOTE: We assume either pos ids to have bsz == 1 (broadcastable) or bsz == effective bsz (input_shape[0])
+ buffered_token_type_ids = self.token_type_ids.expand(position_ids.shape[0], -1)
+ buffered_token_type_ids = torch.gather(buffered_token_type_ids, dim=1, index=position_ids)
+ token_type_ids = buffered_token_type_ids.expand(batch_size, seq_length)
+ else:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.word_embeddings(input_ids)
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
+ embeddings = inputs_embeds + token_type_embeddings
+
+ position_embeddings = self.position_embeddings(position_ids)
+ embeddings = embeddings + position_embeddings
+
+ embeddings = self.LayerNorm(embeddings)
+ embeddings = self.dropout(embeddings)
+ return embeddings
+
+ @staticmethod
+ def create_position_ids_from_inputs_embeds(inputs_embeds, padding_idx):
+ """
+ We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.
+
+ Args:
+ inputs_embeds: torch.Tensor
+
+ Returns: torch.Tensor
+ """
+ input_shape = inputs_embeds.size()[:-1]
+ sequence_length = input_shape[1]
+
+ position_ids = torch.arange(
+ padding_idx + 1, sequence_length + padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
+ )
+ return position_ids.unsqueeze(0).expand(input_shape)
+
+ @staticmethod
+ def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
+ """
+ Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
+ are ignored. This is modified from fairseq's `utils.make_positions`.
+
+ Args:
+ x: torch.Tensor x:
+
+ Returns: torch.Tensor
+ """
+ # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
+ mask = input_ids.ne(padding_idx).int()
+ incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
+ return incremental_indices.long() + padding_idx
+
+
+@auto_docstring
+class BridgeTowerPreTrainedModel(PreTrainedModel):
+ config: BridgeTowerConfig
+ base_model_prefix = "bridgetower"
+ input_modalities = ("image", "text")
+ supports_gradient_checkpointing = False
+ _no_split_modules = ["BridgeTowerSelfAttention", "BridgeTowerResidualAttention"]
+ _skip_keys_device_placement = "past_key_values"
+ _can_record_outputs = {
+ "hidden_states": BridgeTowerTextLayer,
+ "attentions": BridgeTowerSelfAttention,
+ "cross_attentions": BridgeTowerCrossAttention,
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module: nn.Module):
+ std = self.config.initializer_factor
+ if isinstance(module, BridgeTowerVisionTransformer):
+ proj_std = (self.config.hidden_size**-0.5) * ((2 * self.config.num_hidden_layers) ** -0.5)
+ attn_std = self.config.hidden_size**-0.5
+ fc_std = (2 * self.config.hidden_size) ** -0.5
+ for block in module.transformer.resblocks:
+ init.normal_(block.attn.in_proj_weight, std=attn_std * std)
+ init.zeros_(block.attn.in_proj_bias)
+ init.normal_(block.attn.out_proj.weight, std=proj_std * std)
+ init.normal_(block.mlp.c_fc.weight, std=fc_std * std)
+ init.normal_(block.mlp.c_proj.weight, std=proj_std * std)
+
+ init.normal_(module.embeddings.class_embedding, std=attn_std * std)
+ init.normal_(module.embeddings.position_embedding.weight, std=attn_std * std)
+ elif isinstance(module, (nn.Linear, nn.Conv2d, nn.Embedding)):
+ init.normal_(module.weight, mean=0.0, std=0.05 * std)
+ elif isinstance(module, nn.LayerNorm):
+ init.zeros_(module.bias)
+ init.ones_(module.weight)
+ elif isinstance(module, BridgeTowerForContrastiveLearning):
+ init.constant_(module.logit_scale, self.config.logit_scale_init_value)
+ elif isinstance(module, BridgeTowerVisionEmbeddings):
+ init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1)))
+ elif isinstance(module, BridgeTowerTextEmbeddings):
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+ init.zeros_(module.token_type_ids)
+
+ if isinstance(module, (nn.Linear, BridgeTowerMLMHead)) and module.bias is not None:
+ init.zeros_(module.bias)
+
+
+class BridgeTowerVisionModel(BridgeTowerPreTrainedModel):
+ config: BridgeTowerVisionConfig
+ input_modalities = ("image",)
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.visual = BridgeTowerVisionTransformer(config)
+ self.post_init()
+
+ @property
+ def dtype(self):
+ return self.visual.embeddings.patch_embedding.weight.dtype
+
+ def forward(self, image, image_mask=None, interpolate_pos_encoding=False, **kwargs):
+ return self.visual(image.type(self.dtype), image_mask, interpolate_pos_encoding)
+
+
+@auto_docstring(
+ custom_intro="""
+ The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+ cross-attention is added between the self-attention layers, following the architecture described in *Attention is
+ all you need*_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
+ Kaiser and Illia Polosukhin.
+
+ To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set
+ to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and
+ `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.
+
+ .. _*Attention is all you need*: https://huggingface.co/papers/1706.03762
+ """
+)
+class BridgeTowerTextModel(BridgeTowerPreTrainedModel):
+ config: BridgeTowerTextConfig
+ input_modalities = ("text",)
+
+ def __init__(self, config, add_pooling_layer=True):
+ r"""
+ add_pooling_layer (bool, *optional*, defaults to `True`):
+ Whether to add a pooling layer
+ """
+ super().__init__(config)
+ self.config = config
+ self.gradient_checkpointing = False
+
+ self.embeddings = BridgeTowerTextEmbeddings(config)
+ self.encoder = BridgeTowerTextEncoder(config)
+
+ self.pooler = BridgeTowerPooler(config) if add_pooling_layer else None
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embeddings.word_embeddings
+
+ def set_input_embeddings(self, value):
+ self.embeddings.word_embeddings = value
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ # NOTE: bridgetower with its multimodality has a more complicated scheme making records harder
+ # for now we skip the copies from bert but stay close to the original
+ # copied from transformers.models.bert.modeling_bert.BertModel.forward
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPoolingAndCrossAttentions:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if not self.config.is_decoder:
+ use_cache = False
+
+ if use_cache and past_key_values is None:
+ past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+
+ past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+
+ embedding_output = self.embeddings(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ past_key_values_length=past_key_values_length,
+ )
+
+ attention_mask, encoder_attention_mask = self._create_attention_masks(
+ attention_mask=attention_mask,
+ encoder_attention_mask=encoder_attention_mask,
+ embedding_output=embedding_output,
+ encoder_hidden_states=encoder_hidden_states,
+ past_key_values=past_key_values,
+ )
+
+ encoder_outputs = self.encoder(
+ embedding_output,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_ids=position_ids,
+ **kwargs,
+ )
+ sequence_output = encoder_outputs[0]
+ pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+ return BaseModelOutputWithPoolingAndCrossAttentions(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ past_key_values=encoder_outputs.past_key_values,
+ )
+
+ # Copied from transformers.models.bert.modeling_bert.BertModel._create_attention_masks
+ def _create_attention_masks(
+ self,
+ attention_mask,
+ encoder_attention_mask,
+ embedding_output,
+ encoder_hidden_states,
+ past_key_values,
+ ):
+ if self.config.is_decoder:
+ attention_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=embedding_output,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ )
+ else:
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=embedding_output,
+ attention_mask=attention_mask,
+ )
+
+ if encoder_attention_mask is not None:
+ encoder_attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=embedding_output,
+ attention_mask=encoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ )
+
+ return attention_mask, encoder_attention_mask
+
+
+@auto_docstring(
+ custom_intro="""
+ The bare BridgeTower Model transformer outputting BridgeTowerModelOutput object without any specific head on
+ """
+)
+class BridgeTowerModel(BridgeTowerPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.config = config
+ vision_config = config.vision_config
+ text_config = config.text_config
+
+ if config.share_cross_modal_transformer_layers:
+ self.cross_modal_text_transform = nn.Linear(text_config.hidden_size, config.hidden_size)
+ self.cross_modal_image_transform = nn.Linear(vision_config.hidden_size, config.hidden_size)
+ else:
+ self.cross_modal_text_transform = nn.ModuleList(
+ [nn.Linear(text_config.hidden_size, config.hidden_size) for _ in range(config.num_hidden_layers)]
+ )
+ self.cross_modal_image_transform = nn.ModuleList(
+ [nn.Linear(vision_config.hidden_size, config.hidden_size) for _ in range(config.num_hidden_layers)]
+ )
+
+ self.token_type_embeddings = nn.Embedding(2, config.hidden_size)
+
+ self.vision_model = BridgeTowerVisionModel(vision_config)
+
+ self.text_model = BridgeTowerTextModel(text_config)
+
+ if not vision_config.share_layernorm and config.init_layernorm_from_vision_encoder:
+ for ln in self.vision_model.visual.cross_modal_ln_separate:
+ ln.weight.data = self.vision_model.visual.ln_post.weight.data
+ ln.bias.data = self.vision_model.visual.ln_post.bias.data
+
+ self.cross_modal_image_layers = nn.ModuleList(
+ [BridgeTowerBertCrossLayer(text_config) for _ in range(config.num_hidden_layers)]
+ )
+ self.cross_modal_text_layers = nn.ModuleList(
+ [BridgeTowerBertCrossLayer(text_config) for _ in range(config.num_hidden_layers)]
+ )
+
+ # Class token => Linear => Tanh
+ self.cross_modal_image_pooler = BridgeTowerPooler(config)
+ self.cross_modal_text_pooler = BridgeTowerPooler(config)
+
+ # Initialize BridgeTower Components
+ self.cross_modal_text_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.cross_modal_image_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+ if config.share_link_tower_layers:
+ self.cross_modal_text_link_tower = BridgeTowerLinkTower(config)
+ self.cross_modal_image_link_tower = BridgeTowerLinkTower(config)
+ else:
+ self.cross_modal_text_link_tower = nn.ModuleList(
+ [BridgeTowerLinkTower(config) for _ in range(config.num_hidden_layers - 1)]
+ )
+ self.cross_modal_image_link_tower = nn.ModuleList(
+ [BridgeTowerLinkTower(config) for _ in range(config.num_hidden_layers - 1)]
+ )
+
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.text_model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.text_model.set_input_embeddings(value)
+
+ def _apply_text_transform(self, hidden_states: torch.Tensor, layer_idx: int) -> torch.Tensor:
+ if self.config.share_cross_modal_transformer_layers:
+ return self.cross_modal_text_transform(hidden_states)
+ return self.cross_modal_text_transform[layer_idx](hidden_states)
+
+ def _apply_image_transform(self, hidden_states: torch.Tensor, layer_idx: int) -> torch.Tensor:
+ if self.config.share_cross_modal_transformer_layers:
+ return self.cross_modal_image_transform(hidden_states)
+ return self.cross_modal_image_transform[layer_idx](hidden_states)
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ pixel_mask: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ image_embeds: torch.FloatTensor | None = None,
+ image_token_type_idx: int | None = None,
+ labels: torch.LongTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | BridgeTowerModelOutput:
+ r"""
+ image_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`, *optional*):
+ Optionally, instead of passing `pixel_values`, you can choose to directly pass an embedded representation.
+ This is useful if you want more control over how to convert `pixel_values` into patch embeddings.
+ image_token_type_idx (`int`, *optional*):
+ - The token type ids for images.
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels are currently not supported.
+
+ Examples:
+
+ ```python
+ >>> from transformers import BridgeTowerProcessor, BridgeTowerModel
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+
+ >>> # prepare image and text
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+ >>> text = "hello world"
+ >>> processor = BridgeTowerProcessor.from_pretrained("BridgeTower/bridgetower-base")
+ >>> model = BridgeTowerModel.from_pretrained("BridgeTower/bridgetower-base")
+
+ >>> inputs = processor(image, text, return_tensors="pt")
+ >>> outputs = model(**inputs)
+ >>> outputs.keys()
+ odict_keys(['text_features', 'image_features', 'pooler_output'])
+ ```"""
+ all_hidden_states_text = []
+ all_hidden_states_image = []
+ all_hidden_states_cross = []
+ all_self_attentions = []
+
+ if inputs_embeds is not None and input_ids is None:
+ raise NotImplementedError(
+ "BridgeTowerModel does not use `inputs_embeds`. Make sure to pass in `input_ids` instead."
+ )
+
+ image_token_type_idx = image_token_type_idx or 1
+ input_shape = input_ids.size()
+ text_embeds = self.text_model.embeddings(input_ids=input_ids)
+ all_hidden_states_text.append(text_embeds)
+
+ if attention_mask is None:
+ attention_mask = torch.ones(input_shape, dtype=torch.long, device=input_ids.device)
+ extend_text_masks = self.text_model.get_extended_attention_mask(attention_mask, input_shape).to(
+ input_ids.device
+ )
+
+ # The split_index determines how many layers of the uni-modal encoder are applied before the cross-modal encoder
+ split_index = len(self.text_model.encoder.layer) - self.config.num_hidden_layers + 1
+
+ # Run the first 'split_index' layers of the textual encoder
+ for layer in self.text_model.encoder.layer[:split_index]:
+ text_embeds = layer(text_embeds, extend_text_masks)
+ all_hidden_states_text.append(text_embeds)
+
+ if image_embeds is None:
+ image_embeds = self.vision_model.visual.forward_pre(
+ pixel_values.type(self.vision_model.dtype), interpolate_pos_encoding=interpolate_pos_encoding
+ )
+ else:
+ # Permute as BridgeTowerResidualAttention has batch_first=True
+ image_embeds = image_embeds.permute(1, 0, 2)
+
+ all_hidden_states_image.append(image_embeds)
+
+ # Run the first 'split_index' layers of the visual encoder
+ for block in self.vision_model.visual.transformer.resblocks[:split_index]:
+ image_embeds = block(image_embeds)
+ all_hidden_states_image.append(image_embeds)
+
+ image_embeds_with_ln = self.vision_model.visual.forward_post(image_embeds.type(self.vision_model.dtype))
+
+ # first layer is a special case because we don't have the output from the cross-encoder yet
+ cross_modal_text = self._apply_text_transform(text_embeds, layer_idx=0)
+
+ text_token_type_embeddings = self.token_type_embeddings(
+ torch.zeros(1, dtype=torch.long, device=input_ids.device)
+ ).expand_as(cross_modal_text)
+
+ cross_modal_text = self.cross_modal_text_layernorm(cross_modal_text + text_token_type_embeddings)
+
+ image_embeds_with_ln = self._apply_image_transform(image_embeds_with_ln, layer_idx=0)
+ image_token_type_embeddings = self.token_type_embeddings(
+ torch.full((1,), image_token_type_idx, dtype=torch.long, device=input_ids.device)
+ ).expand_as(image_embeds_with_ln)
+
+ image_embeds_with_ln = image_embeds_with_ln + image_token_type_embeddings
+ cross_modal_image = self.cross_modal_image_layernorm(image_embeds_with_ln)
+
+ pixel_mask = torch.ones(
+ (cross_modal_image.size(0), cross_modal_image.size(1)),
+ dtype=torch.long,
+ device=input_ids.device,
+ )
+ extend_image_masks = self.text_model.get_extended_attention_mask(pixel_mask, pixel_mask.size()).to(
+ input_ids.device
+ )
+
+ layer_outputs_text = self.cross_modal_text_layers[0](
+ cross_modal_text,
+ cross_modal_image,
+ attention_mask=extend_text_masks,
+ encoder_attention_mask=extend_image_masks,
+ )
+ cross_text_features = layer_outputs_text[0]
+
+ layer_outputs_image = self.cross_modal_image_layers[0](
+ cross_modal_image,
+ cross_modal_text,
+ attention_mask=extend_image_masks,
+ encoder_attention_mask=extend_text_masks,
+ )
+ cross_image_features = layer_outputs_image[0]
+
+ all_hidden_states_cross.append((cross_text_features, cross_image_features))
+
+ all_self_attentions.append((layer_outputs_text[1], layer_outputs_image[1]))
+
+ link_layer_index = 0
+
+ # Each of the top 6 layers of the visual and textual encoders ([split_index:]) is connected to each layer of
+ # the cross-modal encoder via bridge layers, which brings bottom-up alignment and fusion to the cross-modal encoder.
+ for i in range(split_index, len(self.text_model.encoder.layer)):
+ text_embeds = self.text_model.encoder.layer[i](text_embeds, extend_text_masks)
+ image_embeds = self.vision_model.visual.transformer.resblocks[i](image_embeds).type(
+ self.vision_model.dtype
+ )
+ image_embeds_with_ln = (
+ self._apply_image_transform(self.vision_model.visual.forward_post(image_embeds), link_layer_index + 1)
+ + image_token_type_embeddings
+ )
+
+ text_link_tower = self.cross_modal_text_link_tower[link_layer_index]
+ image_link_tower = self.cross_modal_image_link_tower[link_layer_index]
+
+ # Bridge layers for textual and visual encoders
+ transformed_text_embeds = self._apply_text_transform(text_embeds, link_layer_index + 1)
+ cross_text_features_ = text_link_tower(
+ transformed_text_embeds + text_token_type_embeddings,
+ cross_text_features,
+ extend_text_masks,
+ )
+ cross_image_features_ = image_link_tower(image_embeds_with_ln, cross_image_features, extend_image_masks)
+
+ # Cross-modal encoder via bridge layers of textual and visual encoders
+ layer_outputs_text = self.cross_modal_text_layers[link_layer_index + 1](
+ cross_text_features_,
+ cross_image_features_,
+ attention_mask=extend_text_masks,
+ encoder_attention_mask=extend_image_masks,
+ )
+ cross_text_features = layer_outputs_text[0]
+
+ layer_outputs_image = self.cross_modal_image_layers[link_layer_index + 1](
+ cross_image_features_,
+ cross_text_features_,
+ attention_mask=extend_image_masks,
+ encoder_attention_mask=extend_text_masks,
+ )
+ cross_image_features = layer_outputs_image[0]
+
+ link_layer_index += 1
+
+ all_hidden_states_text.append(text_embeds)
+ all_hidden_states_image.append(image_embeds)
+ all_hidden_states_cross.append((cross_text_features, cross_image_features))
+
+ all_self_attentions.append((layer_outputs_text[1], layer_outputs_image[1]))
+
+ # Concatenate the cls token of the text and image features to get the final represtation
+ text_features, image_features = cross_text_features, cross_image_features
+ cls_features = self.get_cls_features(text_features, image_features)
+
+ return BridgeTowerModelOutput(
+ text_features=text_features,
+ image_features=image_features,
+ pooler_output=cls_features,
+ hidden_states=(
+ tuple(all_hidden_states_text),
+ tuple(all_hidden_states_image),
+ tuple(all_hidden_states_cross),
+ ),
+ attentions=tuple(all_self_attentions),
+ )
+
+ def get_cls_features(self, text_features, image_features):
+ cls_features_text = self.cross_modal_text_pooler(text_features)
+ cls_features_image = self.cross_modal_image_pooler(image_features)
+ return torch.cat([cls_features_text, cls_features_image], dim=-1)
+
+
+# Copied from transformers.models.vilt.modeling_vilt.ViltPredictionHeadTransform with Vilt->BridgeTower
+class BridgeTowerPredictionHeadTransform(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ if isinstance(config.hidden_act, str):
+ self.transform_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.transform_act_fn = config.hidden_act
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+ def forward(self, hidden_states):
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.transform_act_fn(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states)
+ return hidden_states
+
+
+class BridgeTowerMLMHead(nn.Module):
+ def __init__(self, config, weight=None):
+ super().__init__()
+ self.config = config
+ self.transform = BridgeTowerPredictionHeadTransform(config)
+ self.decoder = nn.Linear(config.hidden_size, config.text_config.vocab_size, bias=False)
+ self.bias = nn.Parameter(torch.zeros(config.text_config.vocab_size))
+ if weight is not None:
+ self.decoder.weight = weight
+
+ def forward(self, x):
+ mlm_score = self.transform(x)
+ mlm_score = self.decoder(mlm_score) + self.bias
+ return mlm_score
+
+
+class BridgeTowerITMHead(nn.Module):
+ def __init__(self, hidden_size):
+ super().__init__()
+ self.fc = nn.Linear(hidden_size, 2)
+
+ def forward(self, x):
+ itm_score = self.fc(x)
+ return itm_score
+
+
+@auto_docstring(
+ custom_intro="""
+ BridgeTower Model with a language modeling head on top as done during pretraining.
+ """
+)
+class BridgeTowerForMaskedLM(BridgeTowerPreTrainedModel):
+ _tied_weights_keys = {"mlm_score.decoder.weight": "bridgetower.text_model.embeddings.word_embeddings.weight"}
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.bridgetower = BridgeTowerModel(config)
+ self.mlm_score = BridgeTowerMLMHead(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.mlm_score.decoder
+
+ def set_output_embeddings(self, new_embeddings):
+ self.mlm_score.decoder = new_embeddings
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ pixel_mask: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ image_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> MaskedLMOutput:
+ r"""
+ image_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`, *optional*):
+ Optionally, instead of passing `pixel_values`, you can choose to directly pass an embedded representation.
+ This is useful if you want more control over how to convert `pixel_values` into patch embeddings.
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
+ config.vocab_size]` (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]`
+
+ Examples:
+
+ ```python
+ >>> from transformers import BridgeTowerProcessor, BridgeTowerForMaskedLM
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+
+ >>> url = "http://images.cocodataset.org/val2017/000000360943.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read())).convert("RGB")
+ >>> text = "a looking out of the window"
+
+ >>> processor = BridgeTowerProcessor.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")
+ >>> model = BridgeTowerForMaskedLM.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")
+
+ >>> # prepare inputs
+ >>> encoding = processor(image, text, return_tensors="pt")
+
+ >>> # forward pass
+ >>> outputs = model(**encoding)
+
+ >>> results = processor.decode(outputs.logits.argmax(dim=-1).squeeze(0).tolist())
+
+ >>> print(results)
+ .a cat looking out of the window.
+ ```"""
+ outputs = self.bridgetower(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ pixel_values=pixel_values,
+ pixel_mask=pixel_mask,
+ inputs_embeds=inputs_embeds,
+ image_embeds=image_embeds,
+ **kwargs,
+ )
+
+ mlm_logits = self.mlm_score(outputs.text_features)
+ masked_lm_loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss() # -100 index = padding token
+
+ labels = labels.to(mlm_logits.device)
+ masked_lm_loss = loss_fct(mlm_logits.view(-1, self.config.text_config.vocab_size), labels.view(-1))
+
+ return MaskedLMOutput(
+ loss=masked_lm_loss,
+ logits=mlm_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ BridgeTower Model transformer with a classifier head on top (a linear layer on top of the final hidden state of the
+ [CLS] token) for image-to-text matching.
+ """
+)
+class BridgeTowerForImageAndTextRetrieval(BridgeTowerPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.bridgetower = BridgeTowerModel(config)
+
+ self.itm_score = BridgeTowerITMHead(config.hidden_size * 2)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ pixel_mask: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ image_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> SequenceClassifierOutput:
+ r"""
+ image_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`, *optional*):
+ Optionally, instead of passing `pixel_values`, you can choose to directly pass an embedded representation.
+ This is useful if you want more control over how to convert `pixel_values` into patch embeddings.
+ labels (`torch.LongTensor` of shape `(batch_size, 1)`, *optional*):
+ Labels for computing the image-text matching loss. 0 means the pairs don't match and 1 means they match.
+ The pairs with 0 will be skipped for calculation.
+
+ Examples:
+
+ ```python
+ >>> from transformers import BridgeTowerProcessor, BridgeTowerForImageAndTextRetrieval
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from PIL import Image
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+ >>> texts = ["An image of two cats chilling on a couch", "A football player scoring a goal"]
+
+ >>> processor = BridgeTowerProcessor.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")
+ >>> model = BridgeTowerForImageAndTextRetrieval.from_pretrained("BridgeTower/bridgetower-base-itm-mlm")
+
+ >>> # forward pass
+ >>> scores = dict()
+ >>> for text in texts:
+ ... # prepare inputs
+ ... encoding = processor(image, text, return_tensors="pt")
+ ... outputs = model(**encoding)
+ ... scores[text] = outputs.logits[0, 1].item()
+ ```"""
+ outputs = self.bridgetower(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ pixel_values=pixel_values,
+ pixel_mask=pixel_mask,
+ inputs_embeds=inputs_embeds,
+ image_embeds=image_embeds,
+ **kwargs,
+ )
+
+ pooler_output = outputs.pooler_output
+
+ logits = self.itm_score(pooler_output)
+
+ itm_loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+
+ labels = labels.to(logits.device)
+ itm_loss = loss_fct(logits, labels)
+
+ return SequenceClassifierOutput(
+ loss=itm_loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+class BridgeTowerContrastiveHead(nn.Module):
+ def __init__(self, hidden_size, embed_size):
+ super().__init__()
+ self.fc = nn.Linear(hidden_size, embed_size)
+
+ def forward(self, x):
+ x = self.fc(x)
+ return x
+
+
+@auto_docstring(
+ custom_intro="""
+ BridgeTower Model with a image-text contrastive head on top computing image-text contrastive loss.
+ """
+)
+class BridgeTowerForContrastiveLearning(BridgeTowerPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.bridgetower = BridgeTowerModel(config)
+
+ self.itc_text_head = BridgeTowerContrastiveHead(config.hidden_size, config.contrastive_hidden_size)
+ self.itc_image_head = BridgeTowerContrastiveHead(config.hidden_size, config.contrastive_hidden_size)
+ self.itc_cross_modal_head = BridgeTowerContrastiveHead(config.hidden_size * 2, config.contrastive_hidden_size)
+
+ self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ pixel_mask: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ image_embeds: torch.FloatTensor | None = None,
+ return_loss: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BridgeTowerContrastiveOutput:
+ r"""
+ image_embeds (`torch.FloatTensor` of shape `(batch_size, num_patches, hidden_size)`, *optional*):
+ Optionally, instead of passing `pixel_values`, you can choose to directly pass an embedded representation.
+ This is useful if you want more control over how to convert `pixel_values` into patch embeddings.
+ return_loss (`bool`, *optional*):
+ Whether or not to return the contrastive loss.
+
+ Examples:
+
+ ```python
+ >>> from transformers import BridgeTowerProcessor, BridgeTowerForContrastiveLearning
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from PIL import Image
+ >>> import torch
+
+ >>> image_urls = [
+ ... "https://farm4.staticflickr.com/3395/3428278415_81c3e27f15_z.jpg",
+ ... "http://images.cocodataset.org/val2017/000000039769.jpg",
+ ... ]
+ >>> texts = ["two dogs in a car", "two cats sleeping on a couch"]
+
+ >>> with httpx.stream("GET", urls[0]) as response:
+ ... image1 = Image.open(BytesIO(response.read()))
+
+ >>> with httpx.stream("GET", urls[1]) as response:
+ ... image2 = Image.open(BytesIO(response.read()))
+
+ >>> images = [image1, image2]
+
+ >>> processor = BridgeTowerProcessor.from_pretrained("BridgeTower/bridgetower-large-itm-mlm-itc")
+ >>> model = BridgeTowerForContrastiveLearning.from_pretrained("BridgeTower/bridgetower-large-itm-mlm-itc")
+
+ >>> inputs = processor(images, texts, padding=True, return_tensors="pt")
+ >>> loss = model(**inputs, return_loss=True).loss
+
+ >>> inputs = processor(images, texts[::-1], padding=True, return_tensors="pt")
+ >>> loss_swapped = model(**inputs, return_loss=True).loss
+
+ >>> print("Loss", round(loss.item(), 4))
+ Loss 0.0019
+
+ >>> print("Loss with swapped images", round(loss_swapped.item(), 4))
+ Loss with swapped images 2.126
+ ```"""
+ kwargs.setdefault("output_hidden_states", True)
+ outputs = self.bridgetower(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ pixel_values=pixel_values,
+ pixel_mask=pixel_mask,
+ inputs_embeds=inputs_embeds,
+ image_embeds=image_embeds,
+ **kwargs,
+ )
+
+ pooler_output = outputs.pooler_output
+ hidden_states_txt, hidden_states_img, hidden_states_cross_modal = outputs.hidden_states
+
+ text_embeds = hidden_states_txt[-1]
+ image_embeds = hidden_states_img[-1]
+
+ image_embeds_with_ln = self.bridgetower.vision_model.visual.forward_post(image_embeds)
+ image_token_type_embeddings = self.bridgetower.token_type_embeddings(
+ torch.full((1,), 1, dtype=torch.long, device=self.bridgetower.token_type_embeddings.weight.device)
+ ).expand_as(image_embeds_with_ln)
+
+ image_embeds = self.bridgetower.cross_modal_image_transform(image_embeds_with_ln) + image_token_type_embeddings
+
+ # normalized features
+ text_embeds = nn.functional.normalize(self.itc_text_head(text_embeds[:, 0, :]), dim=-1, p=2)
+ image_embeds = nn.functional.normalize(self.itc_image_head(image_embeds[:, 0, :]), dim=-1, p=2).to(
+ device=text_embeds.device
+ )
+ cross_embeds = nn.functional.normalize(self.itc_cross_modal_head(pooler_output), dim=-1, p=2).to(
+ device=text_embeds.device
+ )
+
+ logits = torch.stack([text_embeds, image_embeds, cross_embeds], dim=-2)
+
+ logit_scale = self.logit_scale.exp().to(device=text_embeds.device)
+ logits_text_to_image = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
+ logits_text_to_cross = torch.matmul(text_embeds, cross_embeds.t()) * logit_scale
+ logits_image_to_cross = torch.matmul(image_embeds, cross_embeds.t()) * logit_scale
+
+ itc_loss = None
+
+ if return_loss:
+ labels = torch.arange(len(logits), device=logits.device)
+ text_to_image_loss = nn.functional.cross_entropy(logits_text_to_image, labels)
+ text_to_cross_loss = nn.functional.cross_entropy(logits_text_to_cross, labels)
+ image_to_cross_loss = nn.functional.cross_entropy(logits_image_to_cross, labels)
+ itc_loss = (text_to_image_loss + text_to_cross_loss + image_to_cross_loss) / 3.0
+
+ return BridgeTowerContrastiveOutput(
+ loss=itc_loss,
+ logits=logits,
+ text_embeds=text_embeds,
+ image_embeds=image_embeds,
+ cross_embeds=cross_embeds,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "BridgeTowerForContrastiveLearning",
+ "BridgeTowerForImageAndTextRetrieval",
+ "BridgeTowerForMaskedLM",
+ "BridgeTowerModel",
+ "BridgeTowerPreTrainedModel",
+]
diff --git a/third_party/transformers/src/transformers/models/bridgetower/processing_bridgetower.py b/third_party/transformers/src/transformers/models/bridgetower/processing_bridgetower.py
new file mode 100644
index 0000000000000000000000000000000000000000..aa0ea7b4c4daa66bffd3757430a252ae7e60081e
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/bridgetower/processing_bridgetower.py
@@ -0,0 +1,49 @@
+# Copyright 2023 The Intel Labs Team Authors, The Microsoft Research Team Authors and HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""
+Processor class for BridgeTower.
+"""
+
+from ...processing_utils import ProcessingKwargs, ProcessorMixin
+from ...utils import auto_docstring
+
+
+class BridgeTowerProcessorKwargs(ProcessingKwargs, total=False):
+ _defaults = {
+ "text_kwargs": {
+ "add_special_tokens": True,
+ "padding": False,
+ "stride": 0,
+ "return_overflowing_tokens": False,
+ "return_special_tokens_mask": False,
+ "return_offsets_mapping": False,
+ "return_length": False,
+ "verbose": True,
+ },
+ "images_kwargs": {
+ "do_normalize": True,
+ "do_center_crop": True,
+ },
+ }
+
+
+@auto_docstring
+class BridgeTowerProcessor(ProcessorMixin):
+ valid_processor_kwargs = BridgeTowerProcessorKwargs
+
+ def __init__(self, image_processor, tokenizer):
+ super().__init__(image_processor, tokenizer)
+
+
+__all__ = ["BridgeTowerProcessor"]
diff --git a/third_party/transformers/src/transformers/models/clap/__init__.py b/third_party/transformers/src/transformers/models/clap/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6d54ee86aecef2cbe5b9bfdee321a0375d977880
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/clap/__init__.py
@@ -0,0 +1,29 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_clap import *
+ from .feature_extraction_clap import *
+ from .modeling_clap import *
+ from .processing_clap import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/clap/configuration_clap.py b/third_party/transformers/src/transformers/models/clap/configuration_clap.py
new file mode 100644
index 0000000000000000000000000000000000000000..4d04623af395e276b2492f2e7845241d86e20c28
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/clap/configuration_clap.py
@@ -0,0 +1,204 @@
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""CLAP model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="laion/clap-htsat-fused")
+@strict
+class ClapTextConfig(PreTrainedConfig):
+ r"""
+ Examples:
+
+ ```python
+ >>> from transformers import ClapTextConfig, ClapTextModel
+
+ >>> # Initializing a CLAP text configuration
+ >>> configuration = ClapTextConfig()
+
+ >>> # Initializing a model (with random weights) from the configuration
+ >>> model = ClapTextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "clap_text_model"
+ base_config_key = "text_config"
+
+ vocab_size: int = 50265
+ hidden_size: int = 768
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 12
+ intermediate_size: int = 3072
+ hidden_act: str = "gelu"
+ hidden_dropout_prob: float | int = 0.1
+ attention_probs_dropout_prob: float | int = 0.1
+ max_position_embeddings: int = 514
+ type_vocab_size: int = 1
+ initializer_factor: float = 1.0
+ layer_norm_eps: float = 1e-12
+ projection_dim: int = 512
+ pad_token_id: int | None = 1
+ bos_token_id: int | None = 0
+ eos_token_id: int | list[int] | None = 2
+ projection_hidden_act: str = "relu"
+
+
+@auto_docstring(checkpoint="laion/clap-htsat-fused")
+@strict
+class ClapAudioConfig(PreTrainedConfig):
+ r"""
+ window_size (`int`, *optional*, defaults to 8):
+ Image size of the spectrogram
+ spec_size (`int`, *optional*, defaults to 256):
+ Desired input size of the spectrogram that the model supports. It can be different from the output of the
+ `ClapFeatureExtractor`, in which case the input features will be resized. Corresponds to the `image_size`
+ of the audio models.
+ patch_stride (`list`, *optional*, defaults to `[4, 4]`):
+ Patch stride for the audio spectrogram
+ num_classes (`int`, *optional*, defaults to 527):
+ Number of classes used for the head training
+ enable_fusion (`bool`, *optional*, defaults to `False`):
+ Whether or not to enable patch fusion. This is the main contribution of the authors, and should give the
+ best results.
+ fusion_type (`[type]`, *optional*):
+ Fusion type used for the patch fusion.
+ patch_embed_input_channels (`int`, *optional*, defaults to 1):
+ Number of channels used for the input spectrogram
+ flatten_patch_embeds (`bool`, *optional*, defaults to `True`):
+ Whether or not to flatten the patch embeddings
+ patch_embeds_hidden_size (`int`, *optional*, defaults to 96):
+ Hidden size of the patch embeddings. It is used as the number of output channels.
+ enable_patch_layer_norm (`bool`, *optional*, defaults to `True`):
+ Whether or not to enable layer normalization for the patch embeddings
+ aff_block_r (`int`, *optional*, defaults to 4):
+ downsize_ratio used in the AudioFF block
+
+ Example:
+
+ ```python
+ >>> from transformers import ClapAudioConfig, ClapAudioModel
+
+ >>> # Initializing a ClapAudioConfig with laion/clap-htsat-fused style configuration
+ >>> configuration = ClapAudioConfig()
+
+ >>> # Initializing a ClapAudioModel (with random weights) from the laion/clap-htsat-fused style configuration
+ >>> model = ClapAudioModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "clap_audio_model"
+ base_config_key = "audio_config"
+
+ window_size: int = 8
+ num_mel_bins: int = 64
+ spec_size: int = 256
+ hidden_act: str = "gelu"
+ patch_size: int | list[int] | tuple[int, int] = 4
+ patch_stride: int | list[int] | tuple[int, ...] = (4, 4)
+ num_classes: int = 527
+ hidden_size: int = 768
+ projection_dim: int = 512
+ depths: list[int] | tuple[int, ...] = (2, 2, 6, 2)
+ num_attention_heads: list[int] | tuple[int, ...] = (4, 8, 16, 32)
+ enable_fusion: bool = False
+ hidden_dropout_prob: float | int = 0.1
+ fusion_type: str | None = None
+ patch_embed_input_channels: int = 1
+ flatten_patch_embeds: bool = True
+ patch_embeds_hidden_size: int = 96
+ enable_patch_layer_norm: bool = True
+ drop_path_rate: float | int = 0.0
+ attention_probs_dropout_prob: float | int = 0.0
+ qkv_bias: bool = True
+ mlp_ratio: float = 4.0
+ aff_block_r: int = 4
+ num_hidden_layers: int = 4
+ projection_hidden_act: str = "relu"
+ layer_norm_eps: float = 1e-5
+ initializer_factor: float = 1.0
+
+
+@auto_docstring(checkpoint="laion/clap-htsat-fused")
+@strict
+class ClapConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import ClapConfig, ClapModel
+
+ >>> # Initializing a ClapConfig with laion-ai/base style configuration
+ >>> configuration = ClapConfig()
+
+ >>> # Initializing a ClapModel (with random weights) from the laion-ai/base style configuration
+ >>> model = ClapModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+
+ >>> # We can also initialize a ClapConfig from a ClapTextConfig and a ClapAudioConfig
+ >>> from transformers import ClapTextConfig, ClapAudioConfig
+
+ >>> # Initializing a ClapText and ClapAudioConfig configuration
+ >>> config_text = ClapTextConfig()
+ >>> config_audio = ClapAudioConfig()
+
+ >>> config = ClapConfig(text_config=config_text, audio_config=config_audio)
+ ```"""
+
+ model_type = "clap"
+ sub_configs = {"text_config": ClapTextConfig, "audio_config": ClapAudioConfig}
+
+ text_config: dict | PreTrainedConfig | None = None
+ audio_config: dict | PreTrainedConfig | None = None
+ logit_scale_init_value: float = 1 / 0.07
+ projection_dim: int = 512
+ projection_hidden_act: str = "relu"
+ initializer_factor: float = 1.0
+
+ def __post_init__(self, **kwargs):
+ if self.text_config is None:
+ self.text_config = ClapTextConfig()
+ logger.info("`text_config` is `None`. initializing the `ClapTextConfig` with default values.")
+ elif isinstance(self.text_config, dict):
+ self.text_config = ClapTextConfig(**self.text_config)
+
+ if self.audio_config is None:
+ self.audio_config = ClapAudioConfig()
+ logger.info("`audio_config` is `None`. initializing the `ClapAudioConfig` with default values.")
+ elif isinstance(self.audio_config, dict):
+ self.audio_config = ClapAudioConfig(**self.audio_config)
+
+ self.text_config.projection_dim = self.projection_dim
+ self.audio_config.projection_dim = self.projection_dim
+
+ self.text_config.projection_hidden_act = self.projection_hidden_act
+ self.audio_config.projection_hidden_act = self.projection_hidden_act
+ self.hidden_size = self.text_config.hidden_size
+ self.num_hidden_layers = self.text_config.num_hidden_layers + len(self.audio_config.depths)
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["ClapAudioConfig", "ClapConfig", "ClapTextConfig"]
diff --git a/third_party/transformers/src/transformers/models/clap/convert_clap_original_pytorch_to_hf.py b/third_party/transformers/src/transformers/models/clap/convert_clap_original_pytorch_to_hf.py
new file mode 100644
index 0000000000000000000000000000000000000000..379e86fd8bbe5d0e477469fe5c667f6138bad2b1
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/clap/convert_clap_original_pytorch_to_hf.py
@@ -0,0 +1,132 @@
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+
+import argparse
+import re
+
+from laion_clap import CLAP_Module
+
+from transformers import AutoFeatureExtractor, ClapConfig, ClapModel
+
+
+KEYS_TO_MODIFY_MAPPING = {
+ "text_branch": "text_model",
+ "audio_branch": "audio_model.audio_encoder",
+ "attn": "attention.self",
+ "self.proj": "output.dense",
+ "attention.self_mask": "attn_mask",
+ "mlp.fc1": "intermediate.dense",
+ "mlp.fc2": "output.dense",
+ "norm1": "layernorm_before",
+ "norm2": "layernorm_after",
+ "bn0": "batch_norm",
+}
+
+processor = AutoFeatureExtractor.from_pretrained("laion/clap-htsat-unfused", truncation="rand_trunc")
+
+
+def init_clap(checkpoint_path, model_type, enable_fusion=False):
+ model = CLAP_Module(
+ amodel=model_type,
+ enable_fusion=enable_fusion,
+ )
+ model.load_ckpt(checkpoint_path)
+ return model
+
+
+def get_config_from_original(clap_model):
+ audio_config = {
+ "patch_embeds_hidden_size": clap_model.model.audio_branch.embed_dim,
+ "depths": clap_model.model.audio_branch.depths,
+ "hidden_size": clap_model.model.audio_projection[0].in_features,
+ }
+
+ text_config = {"hidden_size": clap_model.model.text_branch.pooler.dense.in_features}
+
+ return ClapConfig(audio_config=audio_config, text_config=text_config)
+
+
+def rename_state_dict(state_dict):
+ model_state_dict = {}
+
+ sequential_layers_pattern = r".*sequential.(\d+).*"
+ text_projection_pattern = r".*_projection.(\d+).*"
+
+ for key, value in state_dict.items():
+ # check if any key needs to be modified
+ for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
+ if key_to_modify in key:
+ key = key.replace(key_to_modify, new_key)
+
+ if re.match(sequential_layers_pattern, key):
+ # replace sequential layers with list
+ sequential_layer = re.match(sequential_layers_pattern, key).group(1)
+
+ key = key.replace(f"sequential.{sequential_layer}.", f"layers.{int(sequential_layer) // 3}.linear.")
+ elif re.match(text_projection_pattern, key):
+ projecton_layer = int(re.match(text_projection_pattern, key).group(1))
+
+ # Because in CLAP they use `nn.Sequential`...
+ transformers_projection_layer = 1 if projecton_layer == 0 else 2
+
+ key = key.replace(f"_projection.{projecton_layer}.", f"_projection.linear{transformers_projection_layer}.")
+
+ if "audio" and "qkv" in key:
+ # split qkv into query key and value
+ mixed_qkv = value
+ qkv_dim = mixed_qkv.size(0) // 3
+
+ query_layer = mixed_qkv[:qkv_dim]
+ key_layer = mixed_qkv[qkv_dim : qkv_dim * 2]
+ value_layer = mixed_qkv[qkv_dim * 2 :]
+
+ model_state_dict[key.replace("qkv", "query")] = query_layer
+ model_state_dict[key.replace("qkv", "key")] = key_layer
+ model_state_dict[key.replace("qkv", "value")] = value_layer
+ else:
+ model_state_dict[key] = value
+
+ return model_state_dict
+
+
+def convert_clap_checkpoint(checkpoint_path, pytorch_dump_folder_path, config_path, model_type, enable_fusion=False):
+ clap_model = init_clap(checkpoint_path, model_type, enable_fusion=enable_fusion)
+
+ clap_model.eval()
+ state_dict = clap_model.model.state_dict()
+ state_dict = rename_state_dict(state_dict)
+
+ transformers_config = get_config_from_original(clap_model)
+ transformers_config.audio_config.enable_fusion = enable_fusion
+ model = ClapModel(transformers_config)
+
+ # ignore the spectrogram embedding layer
+ model.load_state_dict(state_dict, strict=False)
+
+ model.save_pretrained(pytorch_dump_folder_path)
+ transformers_config.save_pretrained(pytorch_dump_folder_path)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
+ parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to fairseq checkpoint")
+ parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
+ parser.add_argument("--enable_fusion", action="store_true", help="Whether to enable fusion or not")
+ parser.add_argument("--model_type", default="HTSAT-tiny", type=str, help="Whether to enable fusion or not")
+ args = parser.parse_args()
+
+ convert_clap_checkpoint(
+ args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path, args.model_type, args.enable_fusion
+ )
diff --git a/third_party/transformers/src/transformers/models/clap/feature_extraction_clap.py b/third_party/transformers/src/transformers/models/clap/feature_extraction_clap.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f0a34d2cf4e39d6910b6c9711a3fbc7e97fa7db
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/clap/feature_extraction_clap.py
@@ -0,0 +1,364 @@
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Feature extractor class for CLAP."""
+
+import copy
+from typing import Any
+
+import numpy as np
+import torch
+
+from ...audio_utils import mel_filter_bank, spectrogram, window_function
+from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
+from ...feature_extraction_utils import BatchFeature
+from ...utils import TensorType, logging
+from ...utils.import_utils import requires
+
+
+logger = logging.get_logger(__name__)
+
+
+@requires(backends=("torch",))
+class ClapFeatureExtractor(SequenceFeatureExtractor):
+ r"""
+ Constructs a CLAP feature extractor.
+
+ This feature extractor inherits from [`~feature_extraction_sequence_utils.SequenceFeatureExtractor`] which contains
+ most of the main methods. Users should refer to this superclass for more information regarding those methods.
+
+ This class extracts mel-filter bank features from raw speech using a custom numpy implementation of the *Short Time
+ Fourier Transform* (STFT) which should match pytorch's `torch.stft` equivalent.
+
+ Args:
+ feature_size (`int`, *optional*, defaults to 64):
+ The feature dimension of the extracted Mel spectrograms. This corresponds to the number of mel filters
+ (`n_mels`).
+ sampling_rate (`int`, *optional*, defaults to 48000):
+ The sampling rate at which the audio files should be digitalized expressed in hertz (Hz). This only serves
+ to warn users if the audio fed to the feature extractor does not have the same sampling rate.
+ hop_length (`int`,*optional*, defaults to 480):
+ Length of the overlapping windows for the STFT used to obtain the Mel Spectrogram. The audio will be split
+ in smaller `frames` with a step of `hop_length` between each frame.
+ max_length_s (`int`, *optional*, defaults to 10):
+ The maximum input length of the model in seconds. This is used to pad the audio.
+ fft_window_size (`int`, *optional*, defaults to 1024):
+ Size of the window (in samples) on which the Fourier transform is applied. This controls the frequency
+ resolution of the spectrogram. 400 means that the fourier transform is computed on windows of 400 samples.
+ padding_value (`float`, *optional*, defaults to 0.0):
+ Padding value used to pad the audio. Should correspond to silences.
+ return_attention_mask (`bool`, *optional*, defaults to `False`):
+ Whether or not the model should return the attention masks corresponding to the input.
+ frequency_min (`float`, *optional*, defaults to 0):
+ The lowest frequency of interest. The STFT will not be computed for values below this.
+ frequency_max (`float`, *optional*, defaults to 14000):
+ The highest frequency of interest. The STFT will not be computed for values above this.
+ top_db (`float`, *optional*):
+ The highest decibel value used to convert the mel spectrogram to the log scale. For more details see the
+ `audio_utils.power_to_db` function
+ truncation (`str`, *optional*, defaults to `"fusion"`):
+ Truncation pattern for long audio inputs. Two patterns are available:
+ - `fusion` will use `_random_mel_fusion`, which stacks 3 random crops from the mel spectrogram and a
+ downsampled version of the entire mel spectrogram.
+ If `config.fusion` is set to True, shorter audios also need to return 4 mels, which will just be a copy
+ of the original mel obtained from the padded audio.
+ - `rand_trunc` will select a random crop of the mel spectrogram.
+ padding (`str`, *optional*, defaults to `"repeatpad"`):
+ Padding pattern for shorter audio inputs. Three patterns were originally implemented:
+ - `repeatpad`: the audio is repeated, and then padded to fit the `max_length`.
+ - `repeat`: the audio is repeated and then cut to fit the `max_length`
+ - `pad`: the audio is padded.
+ """
+
+ model_input_names = ["input_features", "is_longer"]
+
+ def __init__(
+ self,
+ feature_size=64,
+ sampling_rate=48_000,
+ hop_length=480,
+ max_length_s=10,
+ fft_window_size=1024,
+ padding_value=0.0,
+ return_attention_mask=False, # pad inputs to max length with silence token (zero) and no attention mask
+ frequency_min: float = 0,
+ frequency_max: float = 14_000,
+ top_db: int | None = None,
+ truncation: str = "fusion",
+ padding: str = "repeatpad",
+ **kwargs,
+ ):
+ super().__init__(
+ feature_size=feature_size,
+ sampling_rate=sampling_rate,
+ padding_value=padding_value,
+ return_attention_mask=return_attention_mask,
+ **kwargs,
+ )
+ self.top_db = top_db
+ self.truncation = truncation
+ self.padding = padding
+ self.fft_window_size = fft_window_size
+ self.nb_frequency_bins = (fft_window_size >> 1) + 1
+ self.hop_length = hop_length
+ self.max_length_s = max_length_s
+ self.nb_max_samples = max_length_s * sampling_rate
+ self.sampling_rate = sampling_rate
+ self.frequency_min = frequency_min
+ self.frequency_max = frequency_max
+ self.mel_filters = mel_filter_bank(
+ num_frequency_bins=self.nb_frequency_bins,
+ num_mel_filters=feature_size,
+ min_frequency=frequency_min,
+ max_frequency=frequency_max,
+ sampling_rate=sampling_rate,
+ norm=None,
+ mel_scale="htk",
+ )
+ self.mel_filters_slaney = mel_filter_bank(
+ num_frequency_bins=self.nb_frequency_bins,
+ num_mel_filters=feature_size,
+ min_frequency=frequency_min,
+ max_frequency=frequency_max,
+ sampling_rate=sampling_rate,
+ norm="slaney",
+ mel_scale="slaney",
+ )
+
+ def to_dict(self) -> dict[str, Any]:
+ """
+ Serializes this instance to a Python dictionary.
+
+ Returns:
+ `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance, except for the
+ mel filter banks, which do not need to be saved or printed as they are too long.
+ """
+ output = copy.deepcopy(self.__dict__)
+ output["feature_extractor_type"] = self.__class__.__name__
+ if "mel_filters" in output:
+ del output["mel_filters"]
+ if "mel_filters_slaney" in output:
+ del output["mel_filters_slaney"]
+ return output
+
+ def _np_extract_fbank_features(self, waveform: np.ndarray, mel_filters: np.ndarray | None = None) -> np.ndarray:
+ """
+ Compute the log-mel spectrogram of the provided `waveform` using the Hann window. In CLAP, two different filter
+ banks are used depending on the truncation pattern:
+ - `self.mel_filters`: they correspond to the default parameters of `torchaudio` which can be obtained from
+ calling `torchaudio.transforms.MelSpectrogram().mel_scale.fb`. These filters are used when `truncation`
+ is set to `"fusion"`.
+ - `self.mel_filteres_slaney` : they correspond to the default parameters of `librosa` which used
+ `librosa.filters.mel` when computing the mel spectrogram. These filters were only used in the original
+ implementation when the truncation mode is not `"fusion"`.
+ """
+ log_mel_spectrogram = spectrogram(
+ waveform,
+ window_function(self.fft_window_size, "hann"),
+ frame_length=self.fft_window_size,
+ hop_length=self.hop_length,
+ power=2.0,
+ mel_filters=mel_filters,
+ log_mel="dB",
+ )
+ return log_mel_spectrogram.T
+
+ def _random_mel_fusion(self, mel, total_frames, chunk_frames):
+ ranges = np.array_split(list(range(0, total_frames - chunk_frames + 1)), 3)
+ if len(ranges[1]) == 0:
+ # if the audio is too short, we just use the first chunk
+ ranges[1] = [0]
+ if len(ranges[2]) == 0:
+ # if the audio is too short, we just use the first chunk
+ ranges[2] = [0]
+ # randomly choose index for each part
+ idx_front = np.random.choice(ranges[0])
+ idx_middle = np.random.choice(ranges[1])
+ idx_back = np.random.choice(ranges[2])
+
+ mel_chunk_front = mel[idx_front : idx_front + chunk_frames, :]
+ mel_chunk_middle = mel[idx_middle : idx_middle + chunk_frames, :]
+ mel_chunk_back = mel[idx_back : idx_back + chunk_frames, :]
+
+ mel = torch.tensor(mel[None, None, :])
+ mel_shrink = torch.nn.functional.interpolate(
+ mel, size=[chunk_frames, 64], mode="bilinear", align_corners=False
+ )
+ mel_shrink = mel_shrink[0][0].numpy()
+ mel_fusion = np.stack([mel_shrink, mel_chunk_front, mel_chunk_middle, mel_chunk_back], axis=0)
+ return mel_fusion
+
+ def _get_input_mel(self, waveform: np.ndarray, max_length, truncation, padding) -> np.ndarray:
+ """
+ Extracts the mel spectrogram and prepares it for the mode based on the `truncation` and `padding` arguments.
+ Four different path are possible:
+ - `truncation="fusion"` and the length of the waveform is greater than the max length: the mel spectrogram
+ will be computed on the entire audio. 3 random crops and a dowsampled version of the full mel spectrogram
+ are then stacked together. They will later be used for `feature_fusion`.
+ - `truncation="rand_trunc"` and the length of the waveform is smaller than the max length: the audio is
+ padded based on `padding`.
+ - `truncation="fusion"` and the length of the waveform is smaller than the max length: the audio is padded
+ based on `padding`, and is repeated `4` times.
+ - `truncation="rand_trunc"` and the length of the waveform is greater than the max length: the mel
+ spectrogram will be computed on a random crop of the waveform.
+
+ """
+ if waveform.shape[0] > max_length:
+ if truncation == "rand_trunc":
+ longer = True
+ # random crop to max_length (for compatibility) -> this should be handled by self.pad
+ overflow = len(waveform) - max_length
+ idx = np.random.randint(0, overflow + 1)
+ waveform = waveform[idx : idx + max_length]
+ input_mel = self._np_extract_fbank_features(waveform, self.mel_filters_slaney)[None, :]
+ elif truncation == "fusion":
+ mel = self._np_extract_fbank_features(waveform, self.mel_filters)
+ chunk_frames = max_length // self.hop_length + 1 # the +1 related to how the spectrogram is computed
+ total_frames = mel.shape[0]
+ if chunk_frames == total_frames:
+ # there is a corner case where the audio length is larger than max_length but smaller than max_length+hop_length.
+ # In this case, we just use the whole audio.
+ input_mel = np.stack([mel, mel, mel, mel], axis=0)
+ longer = False
+ else:
+ input_mel = self._random_mel_fusion(mel, total_frames, chunk_frames)
+ longer = True
+ else:
+ raise NotImplementedError(f"data_truncating {truncation} not implemented")
+
+ else:
+ longer = False
+ # only use repeat as a new possible value for padding. you repeat the audio before applying the usual max_length padding
+ if waveform.shape[0] < max_length:
+ if padding == "repeat":
+ n_repeat = int(max_length / len(waveform))
+ waveform = np.tile(waveform, n_repeat + 1)[:max_length]
+ if padding == "repeatpad":
+ n_repeat = int(max_length / len(waveform))
+ waveform = np.tile(waveform, n_repeat)
+ waveform = np.pad(waveform, (0, max_length - waveform.shape[0]), mode="constant", constant_values=0)
+
+ if truncation == "fusion":
+ input_mel = self._np_extract_fbank_features(waveform, self.mel_filters)
+ input_mel = np.stack([input_mel, input_mel, input_mel, input_mel], axis=0)
+ else:
+ input_mel = self._np_extract_fbank_features(waveform, self.mel_filters_slaney)[None, :]
+
+ return input_mel, longer
+
+ def __call__(
+ self,
+ raw_speech: np.ndarray | list[float] | list[np.ndarray] | list[list[float]],
+ truncation: str | None = None,
+ padding: str | None = None,
+ max_length: int | None = None,
+ sampling_rate: int | None = None,
+ return_tensors: str | TensorType | None = None,
+ **kwargs,
+ ) -> BatchFeature:
+ """
+ Main method to featurize and prepare for the model one or several sequence(s).
+
+ Args:
+ raw_speech (`np.ndarray`, `list[float]`, `list[np.ndarray]`, `list[list[float]]`):
+ The sequence or batch of sequences to be padded. Each sequence can be a numpy array, a list of float
+ values, a list of numpy arrays or a list of list of float values. Must be mono channel audio, not
+ stereo, i.e. single float per timestep.
+ truncation (`str`, *optional*):
+ Truncation pattern for long audio inputs. Two patterns are available:
+ - `fusion` will use `_random_mel_fusion`, which stacks 3 random crops from the mel spectrogram and
+ a downsampled version of the entire mel spectrogram.
+ If `config.fusion` is set to True, shorter audios also need to return 4 mels, which will just be a
+ copy of the original mel obtained from the padded audio.
+ - `rand_trunc` will select a random crop of the mel spectrogram.
+ padding (`str`, *optional*):
+ Padding pattern for shorter audio inputs. Three patterns were originally implemented:
+ - `repeatpad`: the audio is repeated, and then padded to fit the `max_length`.
+ - `repeat`: the audio is repeated and then cut to fit the `max_length`
+ - `pad`: the audio is padded.
+ return_tensors (`str` or [`~utils.TensorType`], *optional*):
+ If set, will return tensors instead of list of python integers. Acceptable values are:
+ - `'pt'`: Return PyTorch `torch.np.array` objects.
+ - `'np'`: Return Numpy `np.ndarray` objects.
+ sampling_rate (`int`, *optional*):
+ The sampling rate at which the `raw_speech` input was sampled. It is strongly recommended to pass
+ `sampling_rate` at the forward call to prevent silent errors and allow automatic speech recognition
+ pipeline.
+ """
+ truncation = truncation if truncation is not None else self.truncation
+ padding = padding if padding else self.padding
+
+ if sampling_rate is not None:
+ if sampling_rate != self.sampling_rate:
+ raise ValueError(
+ f"The model corresponding to this feature extractor: {self.__class__.__name__} was trained using a"
+ f" sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input"
+ f" was sampled with {self.sampling_rate} and not {sampling_rate}."
+ )
+ else:
+ logger.warning(
+ f"It is strongly recommended to pass the `sampling_rate` argument to `{self.__class__.__name__}()`. "
+ "Failing to do so can result in silent errors that might be hard to debug."
+ )
+
+ is_batched_numpy = isinstance(raw_speech, np.ndarray) and len(raw_speech.shape) > 1
+ if is_batched_numpy and len(raw_speech.shape) > 2:
+ raise ValueError(f"Only mono-channel audio is supported for input to {self}")
+ is_batched = is_batched_numpy or (
+ isinstance(raw_speech, (list, tuple)) and (isinstance(raw_speech[0], (np.ndarray, tuple, list)))
+ )
+
+ if is_batched:
+ raw_speech = [np.asarray(speech, dtype=np.float64) for speech in raw_speech]
+ elif not is_batched and not isinstance(raw_speech, np.ndarray):
+ raw_speech = np.asarray(raw_speech, dtype=np.float64)
+ elif isinstance(raw_speech, np.ndarray) and raw_speech.dtype is np.dtype(np.float64):
+ raw_speech = raw_speech.astype(np.float64)
+
+ # always return batch
+ if not is_batched:
+ raw_speech = [np.asarray(raw_speech)]
+
+ # convert to mel spectrogram, truncate and pad if needed.
+ padded_inputs = [
+ self._get_input_mel(waveform, max_length if max_length else self.nb_max_samples, truncation, padding)
+ for waveform in raw_speech
+ ]
+
+ input_mel = []
+ is_longer = []
+ for mel, longer in padded_inputs:
+ input_mel.append(mel)
+ is_longer.append(longer)
+
+ if truncation == "fusion" and sum(is_longer) == 0:
+ # if no audio is longer than 10s, then randomly select one audio to be longer
+ rand_idx = np.random.randint(0, len(input_mel))
+ is_longer[rand_idx] = True
+
+ if isinstance(input_mel[0], list):
+ input_mel = [np.asarray(feature, dtype=np.float64) for feature in input_mel]
+
+ # is_longer is a list of bool
+ is_longer = [[longer] for longer in is_longer]
+
+ input_features = {"input_features": input_mel, "is_longer": is_longer}
+ input_features = BatchFeature(input_features)
+
+ if return_tensors is not None:
+ input_features = input_features.convert_to_tensors(return_tensors)
+
+ return input_features
+
+
+__all__ = ["ClapFeatureExtractor"]
diff --git a/third_party/transformers/src/transformers/models/clap/modeling_clap.py b/third_party/transformers/src/transformers/models/clap/modeling_clap.py
new file mode 100644
index 0000000000000000000000000000000000000000..96c540a3424f0733c900edbbba9987c01d6e98ef
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/clap/modeling_clap.py
@@ -0,0 +1,1775 @@
+# Copyright 2023 The LAION-AI Team and The HuggingFace Team. All rights reserved.
+#
+# 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 CLAP model."""
+
+import collections
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPooling,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...pytorch_utils import apply_chunking_to_forward
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_clap import ClapAudioConfig, ClapConfig, ClapTextConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# Adapted from: https://github.com/LAION-AI/CLAP/blob/6ad05a971ba0622f6acee8c41993e0d02bbed639/src/open_clip/utils.py#L191
+def interpolate(hidden_states, ratio):
+ """
+ Interpolate data in time domain. This is used to compensate the resolution reduction in downsampling of a CNN.
+
+ Args:
+ hidden_states (`torch.FloatTensor` of shape (batch_size, time_length, classes_num)):
+ Input hidden states
+ ratio (`int`):
+ The ratio of the length of the output to the length of the input.
+ """
+ (batch_size, time_length, classes_num) = hidden_states.shape
+ upsampled = hidden_states[:, :, None, :].repeat(1, 1, ratio, 1)
+ upsampled = upsampled.reshape(batch_size, time_length * ratio, classes_num)
+ return upsampled
+
+
+# Adapted from https://github.com/LAION-AI/CLAP/blob/6ad05a971ba0622f6acee8c41993e0d02bbed639/src/open_clip/htsat.py#L249
+def window_partition(hidden_states, window_size):
+ """
+ Returns the resized hidden states. The output shape should be `(batch_size * num_windows, window_size, window_size,
+ num_channels)`
+
+ Args:
+ hidden_states (`torch.FloatTensor` of shape `(batch_size, height, width, num_channels)`):
+ Input hidden states
+ window_size (`int`):
+ Window size
+ """
+ batch_size, height, width, num_channels = hidden_states.shape
+
+ hidden_states = hidden_states.view(
+ batch_size, height // window_size, window_size, width // window_size, window_size, num_channels
+ )
+ windows = hidden_states.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, num_channels)
+ return windows
+
+
+# Adapted from https://github.com/LAION-AI/CLAP/blob/6ad05a971ba0622f6acee8c41993e0d02bbed639/src/open_clip/htsat.py#L263
+def window_reverse(windows, window_size, height, width):
+ """
+ Merges windows to produce higher resolution features.
+ Args:
+ windows (`torch.FloatTensor` of shape `(num_windows * batch_size, window_size, window_size, num_channels)`):
+ Input windows
+ window_size (`int`):
+ Window size
+ height (`int`):
+ Height of the resized audio
+ width (`int`):
+ Width of the resized audio
+ """
+ num_channels = windows.shape[-1]
+ windows = windows.view(-1, height // window_size, width // window_size, window_size, window_size, num_channels)
+ windows = windows.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, height, width, num_channels)
+ return windows
+
+
+# contrastive loss function, adapted from
+# https://sachinruk.github.io/blog/pytorch/pytorch%20lightning/loss%20function/gpu/2021/03/07/CLIP.html#CLIP-loss-function
+def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:
+ labels = torch.arange(len(logits), device=logits.device)
+ return nn.functional.cross_entropy(logits, labels)
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for text model's outputs that also contains a pooling of the last hidden states.
+ """
+)
+# Copied from transformers.models.clip.modeling_clip.CLIPTextModelOutput with CLIP->Clap
+class ClapTextModelOutput(ModelOutput):
+ r"""
+ text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
+ The text embeddings obtained by applying the projection layer to the pooler_output.
+ """
+
+ text_embeds: torch.FloatTensor | None = None
+ last_hidden_state: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ ClapAudio model output to mimic the output of the original implementation.
+ """
+)
+class ClapAudioModelOutput(ModelOutput):
+ r"""
+ audio_embeds (`torch.FloatTensor` of shape `(batch_size, hidden_size)`):
+ The Audio embeddings obtained by applying the projection layer to the pooler_output.
+ """
+
+ audio_embeds: torch.FloatTensor | None = None
+ last_hidden_state: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+
+
+@dataclass
+@auto_docstring
+# Copied from transformers.models.clip.modeling_clip.CLIPOutput with CLIP->Clap, vision->audio, Vision->Audio, image->audio
+class ClapOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`):
+ Contrastive loss for audio-text similarity.
+ logits_per_audio (`torch.FloatTensor` of shape `(audio_batch_size, text_batch_size)`):
+ The scaled dot product scores between `audio_embeds` and `text_embeds`. This represents the audio-text
+ similarity scores.
+ logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, audio_batch_size)`):
+ The scaled dot product scores between `text_embeds` and `audio_embeds`. This represents the text-audio
+ similarity scores.
+ text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):
+ The text embeddings obtained by applying the projection layer to the pooled output of [`ClapTextModel`].
+ audio_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):
+ The audio embeddings obtained by applying the projection layer to the pooled output of [`ClapAudioModel`].
+ text_model_output (`BaseModelOutputWithPooling`):
+ The output of the [`ClapTextModel`].
+ audio_model_output (`BaseModelOutputWithPooling`):
+ The output of the [`ClapAudioModel`].
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits_per_audio: torch.FloatTensor | None = None
+ logits_per_text: torch.FloatTensor | None = None
+ text_embeds: torch.FloatTensor | None = None
+ audio_embeds: torch.FloatTensor | None = None
+ text_model_output: BaseModelOutputWithPooling = None
+ audio_model_output: BaseModelOutputWithPooling = None
+
+ def to_tuple(self) -> tuple[Any]:
+ return tuple(v.to_tuple() if isinstance(v, ModelOutput) else v for v in self.values())
+
+
+# Adapted from transformers.models.swin.modeling_swin.SwinDropPath
+class ClapDropPath(nn.Module):
+ """
+ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). This is a slightly
+ refactored version of the `SwinDropPath` implementation.
+ """
+
+ def __init__(self, drop_prob=None):
+ super().__init__()
+ self.drop_prob = drop_prob
+
+ def forward(self, hidden_states):
+ if self.drop_prob == 0.0 or not self.training:
+ return hidden_states
+
+ keep_prob = 1 - self.drop_prob
+ # work with diff dim tensors, not just 2D ConvNets
+ shape = (hidden_states.shape[0],) + (1,) * (hidden_states.ndim - 1)
+
+ random_tensor = keep_prob + torch.rand(shape, dtype=hidden_states.dtype, device=hidden_states.device)
+ random_tensor.floor_() # binarize
+ output = hidden_states.div(keep_prob) * random_tensor
+ return output
+
+
+# Adapted from https://github.com/LAION-AI/CLAP/blob/6ad05a971ba0622f6acee8c41993e0d02bbed639/src/open_clip/feature_fusion.py#L133
+class ClapAudioAFFBlock(nn.Module):
+ r"""
+ ATTENTIONAL FEATURE FUSION Block from CLAP, since in CLAP we are always in 2D mode, it is not needed to implement
+ the 1D version.
+ """
+
+ def __init__(self, config: ClapAudioConfig):
+ super().__init__()
+ channels = config.patch_embeds_hidden_size
+ downsize_ratio = config.aff_block_r
+ inter_channels = int(channels // downsize_ratio)
+
+ self.local_att = nn.Sequential(
+ nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+ nn.BatchNorm2d(inter_channels),
+ nn.ReLU(inplace=True),
+ nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+ nn.BatchNorm2d(channels),
+ )
+ self.global_att = nn.Sequential(
+ nn.AdaptiveAvgPool2d(1),
+ nn.Conv2d(channels, inter_channels, kernel_size=1, stride=1, padding=0),
+ nn.BatchNorm2d(inter_channels),
+ nn.ReLU(inplace=True),
+ nn.Conv2d(inter_channels, channels, kernel_size=1, stride=1, padding=0),
+ nn.BatchNorm2d(channels),
+ )
+
+ self.sigmoid = nn.Sigmoid()
+
+ def forward(self, hidden_states, residual):
+ attention_input = hidden_states + residual
+
+ fused_layer_output = self.local_att(attention_input) + self.global_att(attention_input)
+ fused_layer_output = self.sigmoid(fused_layer_output)
+
+ output = 2 * hidden_states * fused_layer_output + 2 * residual * (1 - fused_layer_output)
+ return output
+
+
+class ClapAudioPatchEmbed(nn.Module):
+ """
+ This module converts the hidden states reshaped as an image to patch embeddings ready to be passed to the
+ Transformer block.
+ """
+
+ def __init__(self, config: ClapAudioConfig):
+ super().__init__()
+ img_size = (config.spec_size, config.spec_size) if isinstance(config.spec_size, int) else config.spec_size
+ patch_size = (
+ (config.patch_size, config.patch_size) if isinstance(config.patch_size, int) else config.patch_size
+ )
+ patch_stride = (
+ (config.patch_stride, config.patch_stride) if isinstance(config.patch_stride, int) else config.patch_stride
+ )
+
+ self.img_size = img_size
+ self.patch_stride = patch_stride
+
+ self.grid_size = (img_size[0] // patch_stride[0], img_size[1] // patch_stride[1])
+ self.num_patches = self.grid_size[0] * self.grid_size[1]
+
+ self.flatten = config.flatten_patch_embeds
+ self.enable_fusion = config.enable_fusion
+
+ padding = ((patch_size[0] - patch_stride[0]) // 2, (patch_size[1] - patch_stride[1]) // 2)
+
+ scale_factor = 4 if self.enable_fusion and config.fusion_type == "channel_map" else 1
+
+ self.proj = nn.Conv2d(
+ config.patch_embed_input_channels * scale_factor,
+ config.patch_embeds_hidden_size,
+ kernel_size=patch_size,
+ stride=patch_stride,
+ padding=padding,
+ )
+
+ self.norm = nn.LayerNorm(config.patch_embeds_hidden_size) if config.enable_patch_layer_norm else nn.Identity()
+ if self.enable_fusion:
+ self.fusion_model = ClapAudioAFFBlock(config)
+ self.mel_conv2d = nn.Conv2d(
+ config.patch_embed_input_channels,
+ config.patch_embeds_hidden_size,
+ kernel_size=(patch_size[0], patch_size[1] * 3),
+ stride=(patch_stride[0], patch_stride[1] * 3),
+ padding=padding,
+ )
+
+ def forward(self, hidden_states, is_longer_idx=None):
+ if self.enable_fusion:
+ # retrieve the last mel as we have transposed the input
+ global_hidden_states = hidden_states[:, 0:1, :, :]
+
+ # global processing
+ batch_size, num_channels, height, width = global_hidden_states.shape
+
+ if height != self.img_size[0] or width != self.img_size[1]:
+ raise ValueError(
+ f"Input audio size ({height}*{width}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+ )
+
+ global_hidden_states = self.proj(global_hidden_states)
+ output_width = global_hidden_states.size(-1)
+ if len(is_longer_idx) > 0:
+ # local processing
+ local_hidden_states = hidden_states[is_longer_idx, 1:, :, :].contiguous()
+ batch_size, num_channels, height, width = local_hidden_states.shape
+ local_hidden_states = local_hidden_states.view(batch_size * num_channels, 1, height, width)
+
+ local_hidden_states = self.mel_conv2d(local_hidden_states)
+
+ _, features, height, width = local_hidden_states.shape
+ local_hidden_states = local_hidden_states.view(batch_size, num_channels, features, height, width)
+ local_hidden_states = local_hidden_states.permute((0, 2, 3, 1, 4)).contiguous().flatten(3)
+
+ local_width = local_hidden_states.size(-1)
+ local_hidden_states = torch.nn.functional.pad(
+ local_hidden_states, (0, output_width - local_width), "constant", 0
+ )
+
+ global_hidden_states[is_longer_idx] = self.fusion_model(
+ global_hidden_states[is_longer_idx], local_hidden_states
+ )
+ hidden_states = global_hidden_states
+ else:
+ _, _, height, width = hidden_states.shape
+ if height != self.img_size[0] or width != self.img_size[1]:
+ raise ValueError(
+ f"Input audio size ({height}*{width}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+ )
+ hidden_states = self.proj(hidden_states)
+
+ if self.flatten:
+ hidden_states = hidden_states.flatten(2).transpose(1, 2)
+ hidden_states = self.norm(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.swin.modeling_swin.SwinSelfAttention with Swin->ClapAudio
+class ClapAudioSelfAttention(nn.Module):
+ def __init__(self, config, dim, num_heads, window_size):
+ super().__init__()
+ if dim % num_heads != 0:
+ raise ValueError(
+ f"The hidden size ({dim}) is not a multiple of the number of attention heads ({num_heads})"
+ )
+
+ self.num_attention_heads = num_heads
+ self.attention_head_size = int(dim / num_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+ self.window_size = (
+ window_size if isinstance(window_size, collections.abc.Iterable) else (window_size, window_size)
+ )
+
+ self.relative_position_bias_table = nn.Parameter(
+ torch.zeros((2 * self.window_size[0] - 1) * (2 * self.window_size[1] - 1), num_heads)
+ )
+
+ self.register_buffer("relative_position_index", self.create_relative_position_index())
+
+ self.query = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
+ self.key = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
+ self.value = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ output_attentions: bool | None = False,
+ ) -> tuple[torch.Tensor]:
+ batch_size, dim, num_channels = hidden_states.shape
+ hidden_shape = (batch_size, dim, -1, self.attention_head_size)
+
+ query_layer = self.query(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_layer = self.key(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_layer = self.value(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+ attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+
+ relative_position_bias = self.relative_position_bias_table[self.relative_position_index.view(-1)]
+ relative_position_bias = relative_position_bias.view(
+ self.window_size[0] * self.window_size[1], self.window_size[0] * self.window_size[1], -1
+ )
+
+ relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()
+ attention_scores = attention_scores + relative_position_bias.unsqueeze(0)
+
+ if attention_mask is not None:
+ # Apply the attention mask is (precomputed for all layers in ClapAudioModel forward() function)
+ mask_shape = attention_mask.shape[0]
+ attention_scores = attention_scores.view(
+ batch_size // mask_shape, mask_shape, self.num_attention_heads, dim, dim
+ )
+ attention_scores = attention_scores + attention_mask.unsqueeze(1).unsqueeze(0)
+ attention_scores = attention_scores.view(-1, self.num_attention_heads, dim, dim)
+
+ # Normalize the attention scores to probabilities.
+ attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+
+ # This is actually dropping out entire tokens to attend to, which might
+ # seem a bit unusual, but is taken from the original Transformer paper.
+ attention_probs = self.dropout(attention_probs)
+
+ context_layer = torch.matmul(attention_probs, value_layer)
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+ context_layer = context_layer.view(new_context_layer_shape)
+
+ outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+
+ return outputs
+
+ def create_relative_position_index(self):
+ # get pair-wise relative position index for each token inside the window
+ coords_h = torch.arange(self.window_size[0])
+ coords_w = torch.arange(self.window_size[1])
+ coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij"))
+ coords_flatten = torch.flatten(coords, 1)
+ relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]
+ relative_coords = relative_coords.permute(1, 2, 0).contiguous()
+ relative_coords[:, :, 0] += self.window_size[0] - 1
+ relative_coords[:, :, 1] += self.window_size[1] - 1
+ relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1
+ relative_position_index = relative_coords.sum(-1)
+ return relative_position_index
+
+
+# Copied from transformers.models.swin.modeling_swin.SwinSelfOutput with Swin->ClapAudio
+class ClapAudioSelfOutput(nn.Module):
+ def __init__(self, config, dim):
+ super().__init__()
+ self.dense = nn.Linear(dim, dim)
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+
+ return hidden_states
+
+
+# Copied from transformers.models.swin.modeling_swin.SwinAttention with Swin->ClapAudio
+class ClapAudioAttention(nn.Module):
+ def __init__(self, config, dim, num_heads, window_size):
+ super().__init__()
+ self.self = ClapAudioSelfAttention(config, dim, num_heads, window_size)
+ self.output = ClapAudioSelfOutput(config, dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ output_attentions: bool | None = False,
+ ) -> tuple[torch.Tensor]:
+ self_outputs = self.self(hidden_states, attention_mask, output_attentions)
+ attention_output = self.output(self_outputs[0], hidden_states)
+ outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
+ return outputs
+
+
+# Copied from transformers.models.swin.modeling_swin.SwinIntermediate with Swin->ClapAudio
+class ClapAudioIntermediate(nn.Module):
+ def __init__(self, config, dim):
+ super().__init__()
+ self.dense = nn.Linear(dim, int(config.mlp_ratio * dim))
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.swin.modeling_swin.SwinOutput with Swin->ClapAudio
+class ClapAudioOutput(nn.Module):
+ def __init__(self, config, dim):
+ super().__init__()
+ self.dense = nn.Linear(int(config.mlp_ratio * dim), dim)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.swin.modeling_swin.SwinLayer with SwinDropPath->ClapDropPath, Swin->ClapAudio
+class ClapAudioLayer(nn.Module):
+ def __init__(self, config, dim, input_resolution, num_heads, drop_path_rate=0.0, shift_size=0):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.shift_size = shift_size
+ self.window_size = config.window_size
+ self.input_resolution = input_resolution
+ self.layernorm_before = nn.LayerNorm(dim, eps=config.layer_norm_eps)
+ self.attention = ClapAudioAttention(config, dim, num_heads, window_size=self.window_size)
+ self.drop_path = ClapDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
+ self.layernorm_after = nn.LayerNorm(dim, eps=config.layer_norm_eps)
+ self.intermediate = ClapAudioIntermediate(config, dim)
+ self.output = ClapAudioOutput(config, dim)
+
+ def set_shift_and_window_size(self, input_resolution):
+ if min(input_resolution) <= self.window_size:
+ # if window size is larger than input resolution, we don't partition windows
+ self.shift_size = torch_int(0)
+ self.window_size = (
+ torch.min(torch.tensor(input_resolution)) if torch.jit.is_tracing() else min(input_resolution)
+ )
+
+ def get_attn_mask(self, height, width, dtype, device):
+ if self.shift_size > 0:
+ # calculate attention mask for SW-MSA
+ img_mask = torch.zeros((1, height, width, 1), dtype=dtype, device=device)
+ height_slices = (
+ slice(0, -self.window_size),
+ slice(-self.window_size, -self.shift_size),
+ slice(-self.shift_size, None),
+ )
+ width_slices = (
+ slice(0, -self.window_size),
+ slice(-self.window_size, -self.shift_size),
+ slice(-self.shift_size, None),
+ )
+ count = 0
+ for height_slice in height_slices:
+ for width_slice in width_slices:
+ img_mask[:, height_slice, width_slice, :] = count
+ count += 1
+
+ mask_windows = window_partition(img_mask, self.window_size)
+ mask_windows = mask_windows.view(-1, self.window_size * self.window_size)
+ attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)
+ attn_mask = attn_mask.masked_fill(attn_mask != 0, -100.0).masked_fill(attn_mask == 0, 0.0)
+ else:
+ attn_mask = None
+ return attn_mask
+
+ def maybe_pad(self, hidden_states, height, width):
+ pad_right = (self.window_size - width % self.window_size) % self.window_size
+ pad_bottom = (self.window_size - height % self.window_size) % self.window_size
+ pad_values = (0, 0, 0, pad_right, 0, pad_bottom)
+ hidden_states = nn.functional.pad(hidden_states, pad_values)
+ return hidden_states, pad_values
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ input_dimensions: tuple[int, int],
+ output_attentions: bool | None = False,
+ always_partition: bool | None = False,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ if not always_partition:
+ self.set_shift_and_window_size(input_dimensions)
+ else:
+ pass
+ height, width = input_dimensions
+ batch_size, _, channels = hidden_states.size()
+ shortcut = hidden_states
+
+ hidden_states = self.layernorm_before(hidden_states)
+
+ hidden_states = hidden_states.view(batch_size, height, width, channels)
+
+ # pad hidden_states to multiples of window size
+ hidden_states, pad_values = self.maybe_pad(hidden_states, height, width)
+
+ _, height_pad, width_pad, _ = hidden_states.shape
+ # cyclic shift
+ if self.shift_size > 0:
+ shifted_hidden_states = torch.roll(hidden_states, shifts=(-self.shift_size, -self.shift_size), dims=(1, 2))
+ else:
+ shifted_hidden_states = hidden_states
+
+ # partition windows
+ hidden_states_windows = window_partition(shifted_hidden_states, self.window_size)
+ hidden_states_windows = hidden_states_windows.view(-1, self.window_size * self.window_size, channels)
+ attn_mask = self.get_attn_mask(
+ height_pad, width_pad, dtype=hidden_states.dtype, device=hidden_states_windows.device
+ )
+
+ attention_outputs = self.attention(hidden_states_windows, attn_mask, output_attentions=output_attentions)
+
+ attention_output = attention_outputs[0]
+
+ attention_windows = attention_output.view(-1, self.window_size, self.window_size, channels)
+ shifted_windows = window_reverse(attention_windows, self.window_size, height_pad, width_pad)
+
+ # reverse cyclic shift
+ if self.shift_size > 0:
+ attention_windows = torch.roll(shifted_windows, shifts=(self.shift_size, self.shift_size), dims=(1, 2))
+ else:
+ attention_windows = shifted_windows
+
+ was_padded = pad_values[3] > 0 or pad_values[5] > 0
+ if was_padded:
+ attention_windows = attention_windows[:, :height, :width, :].contiguous()
+
+ attention_windows = attention_windows.view(batch_size, height * width, channels)
+
+ hidden_states = shortcut + self.drop_path(attention_windows)
+
+ layer_output = self.layernorm_after(hidden_states)
+ layer_output = self.intermediate(layer_output)
+ layer_output = hidden_states + self.output(layer_output)
+
+ layer_outputs = (layer_output, attention_outputs[1]) if output_attentions else (layer_output,)
+ return layer_outputs
+
+
+# Copied from transformers.models.swin.modeling_swin.SwinStage with Swin->ClapAudio
+class ClapAudioStage(GradientCheckpointingLayer):
+ def __init__(self, config, dim, input_resolution, depth, num_heads, drop_path, downsample):
+ super().__init__()
+ self.config = config
+ self.dim = dim
+ self.blocks = nn.ModuleList(
+ [
+ ClapAudioLayer(
+ config=config,
+ dim=dim,
+ input_resolution=input_resolution,
+ num_heads=num_heads,
+ drop_path_rate=drop_path[i],
+ shift_size=0 if (i % 2 == 0) else config.window_size // 2,
+ )
+ for i in range(depth)
+ ]
+ )
+
+ # patch merging layer
+ if downsample is not None:
+ self.downsample = downsample(input_resolution, dim=dim, norm_layer=nn.LayerNorm)
+ else:
+ self.downsample = None
+
+ self.pointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ input_dimensions: tuple[int, int],
+ output_attentions: bool | None = False,
+ always_partition: bool | None = False,
+ ) -> tuple[torch.Tensor]:
+ height, width = input_dimensions
+ for i, layer_module in enumerate(self.blocks):
+ layer_outputs = layer_module(hidden_states, input_dimensions, output_attentions, always_partition)
+
+ hidden_states = layer_outputs[0]
+
+ hidden_states_before_downsampling = hidden_states
+ if self.downsample is not None:
+ height_downsampled, width_downsampled = (height + 1) // 2, (width + 1) // 2
+ output_dimensions = (height, width, height_downsampled, width_downsampled)
+ hidden_states = self.downsample(hidden_states_before_downsampling, input_dimensions)
+ else:
+ output_dimensions = (height, width, height, width)
+
+ stage_outputs = (hidden_states, hidden_states_before_downsampling, output_dimensions)
+
+ if output_attentions:
+ stage_outputs += layer_outputs[1:]
+ return stage_outputs
+
+
+# Copied from transformers.models.swin.modeling_swin.SwinPatchMerging with Swin->ClapAudio
+class ClapAudioPatchMerging(nn.Module):
+ """
+ Patch Merging Layer.
+
+ Args:
+ input_resolution (`tuple[int]`):
+ Resolution of input feature.
+ dim (`int`):
+ Number of input channels.
+ norm_layer (`nn.Module`, *optional*, defaults to `nn.LayerNorm`):
+ Normalization layer class.
+ """
+
+ def __init__(self, input_resolution: tuple[int], dim: int, norm_layer: nn.Module = nn.LayerNorm) -> None:
+ super().__init__()
+ self.input_resolution = input_resolution
+ self.dim = dim
+ self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)
+ self.norm = norm_layer(4 * dim)
+
+ def maybe_pad(self, input_feature, height, width):
+ should_pad = (height % 2 == 1) or (width % 2 == 1)
+ if should_pad:
+ pad_values = (0, 0, 0, width % 2, 0, height % 2)
+ input_feature = nn.functional.pad(input_feature, pad_values)
+
+ return input_feature
+
+ def forward(self, input_feature: torch.Tensor, input_dimensions: tuple[int, int]) -> torch.Tensor:
+ height, width = input_dimensions
+ # `dim` is height * width
+ batch_size, dim, num_channels = input_feature.shape
+
+ input_feature = input_feature.view(batch_size, height, width, num_channels)
+ # pad input to be divisible by width and height, if needed
+ input_feature = self.maybe_pad(input_feature, height, width)
+ # [batch_size, height/2, width/2, num_channels]
+ input_feature_0 = input_feature[:, 0::2, 0::2, :]
+ # [batch_size, height/2, width/2, num_channels]
+ input_feature_1 = input_feature[:, 1::2, 0::2, :]
+ # [batch_size, height/2, width/2, num_channels]
+ input_feature_2 = input_feature[:, 0::2, 1::2, :]
+ # [batch_size, height/2, width/2, num_channels]
+ input_feature_3 = input_feature[:, 1::2, 1::2, :]
+ # batch_size height/2 width/2 4*num_channels
+ input_feature = torch.cat([input_feature_0, input_feature_1, input_feature_2, input_feature_3], -1)
+ input_feature = input_feature.view(batch_size, -1, 4 * num_channels) # batch_size height/2*width/2 4*C
+
+ input_feature = self.norm(input_feature)
+ input_feature = self.reduction(input_feature)
+
+ return input_feature
+
+
+class ClapAudioEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.num_layers = len(config.depths)
+
+ self.config = config
+ self.patch_embed = ClapAudioPatchEmbed(config)
+ self.enable_fusion = config.enable_fusion
+ self.patch_stride = self.patch_embed.patch_stride
+ self.spec_size = config.spec_size
+ self.freq_ratio = config.spec_size // config.num_mel_bins
+
+ self.num_features = int(config.patch_embeds_hidden_size * 2 ** (self.num_layers - 1))
+
+ drop_path_rate = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths), device="cpu")]
+
+ grid_size = self.patch_embed.grid_size
+ self.input_resolutions = [(grid_size[0] // (2**i), grid_size[1] // (2**i)) for i in range(self.num_layers)]
+
+ self.layers = nn.ModuleList(
+ [
+ ClapAudioStage(
+ config=config,
+ dim=int(config.patch_embeds_hidden_size * 2**i_layer),
+ input_resolution=self.input_resolutions[i_layer],
+ depth=config.depths[i_layer],
+ num_heads=config.num_attention_heads[i_layer],
+ drop_path=drop_path_rate[sum(config.depths[:i_layer]) : sum(config.depths[: i_layer + 1])],
+ downsample=ClapAudioPatchMerging if (i_layer < self.num_layers - 1) else None,
+ )
+ for i_layer in range(self.num_layers)
+ ]
+ )
+
+ self.gradient_checkpointing = False
+
+ self.batch_norm = nn.BatchNorm2d(config.num_mel_bins)
+ self.norm = nn.LayerNorm(self.num_features)
+ self.depths = config.depths
+ self.avgpool = nn.AdaptiveAvgPool1d(1)
+
+ def reshape_mel2img(self, normalized_input_features):
+ """
+ The input is 4 normalized log mel spectrograms. It is reshape to the common shape of images. Each channel
+ should represent 1 of the 4 crops of the spectrogram. For more details, refer to the [`ClapFeatureExtractor`].
+ """
+ _, _, time_length, freq_length = normalized_input_features.shape
+
+ spec_width = int(self.spec_size * self.freq_ratio)
+ spec_height = self.spec_size // self.freq_ratio
+
+ if time_length > spec_width or freq_length > spec_height:
+ raise ValueError("the wav size should be less than or equal to the swin input size")
+
+ # to avoid bicubic zero error
+ if time_length < spec_width:
+ normalized_input_features = nn.functional.interpolate(
+ normalized_input_features, (spec_width, freq_length), mode="bicubic", align_corners=True
+ )
+ if freq_length < spec_height:
+ normalized_input_features = nn.functional.interpolate(
+ normalized_input_features, (time_length, spec_height), mode="bicubic", align_corners=True
+ )
+
+ batch, channels, time, freq = normalized_input_features.shape
+
+ # batch_size, channels, spec_width, spec_height --> batch_size, channels, spec_height * freq_ratio, spec_width // freq_ratio
+ normalized_input_features = normalized_input_features.reshape(
+ batch, channels * self.freq_ratio, time // self.freq_ratio, freq
+ )
+ normalized_input_features = normalized_input_features.permute(0, 1, 3, 2).contiguous()
+ normalized_input_features = normalized_input_features.reshape(
+ batch, channels, freq * self.freq_ratio, time // self.freq_ratio
+ )
+
+ return normalized_input_features
+
+ @can_return_tuple
+ def forward(
+ self,
+ input_features,
+ is_longer: torch.FloatTensor | None = None,
+ output_attentions: bool | None = False,
+ output_hidden_states: bool | None = False,
+ output_hidden_states_before_downsampling: bool | None = False,
+ always_partition: bool | None = False,
+ return_dict: bool | None = True,
+ ) -> tuple | ClapAudioModelOutput:
+ # Unique logic so no refactor here yet
+ output_hidden_states = output_hidden_states or self.config.output_hidden_states
+ output_attentions = output_attentions or self.config.output_attentions
+
+ input_features = input_features.transpose(1, 3)
+ normalized_input_features = self.batch_norm(input_features)
+ normalized_input_features = normalized_input_features.transpose(1, 3)
+
+ is_longer_list_idx = None
+ if self.enable_fusion:
+ is_longer_list = is_longer.to(input_features.device)
+ is_longer_list_idx = torch.where(is_longer_list == 1)[0]
+
+ hidden_states = self.reshape_mel2img(normalized_input_features)
+
+ frames_num = hidden_states.shape[2]
+
+ hidden_states = self.patch_embed(hidden_states, is_longer_list_idx)
+
+ all_hidden_states = () if output_hidden_states else None
+ all_reshaped_hidden_states = () if output_hidden_states else None
+ all_self_attentions = () if output_attentions else None
+
+ input_dimensions = self.input_resolutions[0]
+
+ if output_hidden_states:
+ batch_size, _, hidden_size = hidden_states.shape
+ # rearrange batch_size (height width) channels -> batch_size channel height width
+ reshaped_hidden_state = hidden_states.view(batch_size, *input_dimensions, hidden_size)
+ reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
+ all_hidden_states += (hidden_states,)
+ all_reshaped_hidden_states += (reshaped_hidden_state,)
+
+ for i, layer_module in enumerate(self.layers):
+ input_dimensions = self.input_resolutions[i]
+
+ layer_outputs = layer_module(hidden_states, input_dimensions, output_attentions, always_partition)
+
+ hidden_states = layer_outputs[0]
+
+ hidden_states_before_downsampling = layer_outputs[1]
+ output_dimensions = layer_outputs[2]
+
+ input_dimensions = (output_dimensions[-2], output_dimensions[-1])
+
+ if output_hidden_states and output_hidden_states_before_downsampling:
+ batch_size, _, hidden_size = hidden_states_before_downsampling.shape
+ # rearrange batch_size (height width) channels -> batch_size channel height width
+ # here we use the original (not downsampled) height and width
+ reshaped_hidden_state = hidden_states_before_downsampling.view(
+ batch_size, *(output_dimensions[0], output_dimensions[1]), hidden_size
+ )
+ reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
+ all_hidden_states += (hidden_states_before_downsampling,)
+ all_reshaped_hidden_states += (reshaped_hidden_state,)
+ elif output_hidden_states and not output_hidden_states_before_downsampling:
+ batch_size, _, hidden_size = hidden_states.shape
+ # rearrange batch_size (height width) channels -> batch_size channel height width
+ reshaped_hidden_state = hidden_states.view(batch_size, *input_dimensions, hidden_size)
+ reshaped_hidden_state = reshaped_hidden_state.permute(0, 3, 1, 2)
+ all_hidden_states += (hidden_states,)
+ all_reshaped_hidden_states += (reshaped_hidden_state,)
+
+ if output_attentions:
+ all_self_attentions += layer_outputs[3:]
+
+ last_hidden_state = self.norm(hidden_states)
+
+ batch_size, _, n_channels = last_hidden_state.shape
+
+ freq_shape = frames_num // (2 ** (len(self.depths) - 1)) // self.patch_stride[0]
+ temporal_shape = frames_num // (2 ** (len(self.depths) - 1)) // self.patch_stride[1]
+
+ last_hidden_state = (
+ last_hidden_state.permute(0, 2, 1).contiguous().reshape(batch_size, n_channels, freq_shape, temporal_shape)
+ )
+
+ batch_size, n_channels, n_frequencies, n_temp = last_hidden_state.shape
+ # group 2D CNN
+ c_freq_bin = n_frequencies // self.freq_ratio
+ last_hidden_state = last_hidden_state.reshape(
+ batch_size, n_channels, n_frequencies // c_freq_bin, c_freq_bin, n_temp
+ )
+ last_hidden_state = (
+ last_hidden_state.permute(0, 1, 3, 2, 4).contiguous().reshape(batch_size, n_channels, c_freq_bin, -1)
+ )
+ latent_output = self.avgpool(torch.flatten(last_hidden_state, 2))
+ latent_output = torch.flatten(latent_output, 1)
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=last_hidden_state,
+ pooler_output=latent_output,
+ hidden_states=all_reshaped_hidden_states,
+ attentions=all_self_attentions,
+ )
+
+
+class ClapProjectionLayer(nn.Module):
+ def __init__(self, config: ClapAudioConfig | ClapTextConfig):
+ super().__init__()
+ self.config = config
+ hidden_size = config.hidden_size
+ projection_dim = config.projection_dim
+
+ self.linear1 = nn.Linear(hidden_size, projection_dim)
+ self.activation = ACT2FN[config.projection_hidden_act]
+ self.linear2 = nn.Linear(projection_dim, projection_dim)
+
+ def forward(self, hidden_states):
+ hidden_states = self.linear1(hidden_states)
+ hidden_states = self.activation(hidden_states)
+ hidden_states = self.linear2(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.roberta.modeling_roberta.RobertaEmbeddings with Roberta->ClapText, persistent=False->persistent=True
+class ClapTextEmbeddings(nn.Module):
+ """Construct the embeddings from word, position and token_type embeddings."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+ self.register_buffer(
+ "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=True
+ )
+ self.register_buffer(
+ "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=True
+ )
+
+ self.padding_idx = config.pad_token_id
+ self.position_embeddings = nn.Embedding(
+ config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
+ )
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values_length: int = 0,
+ ) -> torch.Tensor:
+ if position_ids is None:
+ if input_ids is not None:
+ # Create the position ids from the input token ids. Any padded tokens remain padded.
+ position_ids = self.create_position_ids_from_input_ids(
+ input_ids, self.padding_idx, past_key_values_length
+ )
+ else:
+ position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds, self.padding_idx)
+
+ if input_ids is not None:
+ input_shape = input_ids.size()
+ else:
+ input_shape = inputs_embeds.size()[:-1]
+
+ batch_size, seq_length = input_shape
+
+ # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
+ # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
+ # issue #5664
+ if token_type_ids is None:
+ if hasattr(self, "token_type_ids"):
+ # NOTE: We assume either pos ids to have bsz == 1 (broadcastable) or bsz == effective bsz (input_shape[0])
+ buffered_token_type_ids = self.token_type_ids.expand(position_ids.shape[0], -1)
+ buffered_token_type_ids = torch.gather(buffered_token_type_ids, dim=1, index=position_ids)
+ token_type_ids = buffered_token_type_ids.expand(batch_size, seq_length)
+ else:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.word_embeddings(input_ids)
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
+ embeddings = inputs_embeds + token_type_embeddings
+
+ position_embeddings = self.position_embeddings(position_ids)
+ embeddings = embeddings + position_embeddings
+
+ embeddings = self.LayerNorm(embeddings)
+ embeddings = self.dropout(embeddings)
+ return embeddings
+
+ @staticmethod
+ def create_position_ids_from_inputs_embeds(inputs_embeds, padding_idx):
+ """
+ We are provided embeddings directly. We cannot infer which are padded so just generate sequential position ids.
+
+ Args:
+ inputs_embeds: torch.Tensor
+
+ Returns: torch.Tensor
+ """
+ input_shape = inputs_embeds.size()[:-1]
+ sequence_length = input_shape[1]
+
+ position_ids = torch.arange(
+ padding_idx + 1, sequence_length + padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
+ )
+ return position_ids.unsqueeze(0).expand(input_shape)
+
+ @staticmethod
+ def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
+ """
+ Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
+ are ignored. This is modified from fairseq's `utils.make_positions`.
+
+ Args:
+ x: torch.Tensor x:
+
+ Returns: torch.Tensor
+ """
+ # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
+ mask = input_ids.ne(padding_idx).int()
+ incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
+ return incremental_indices.long() + padding_idx
+
+
+# Copied from transformers.models.align.modeling_align.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.align.modeling_align.AlignTextSelfAttention with Align->Clap
+class ClapTextSelfAttention(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+ raise ValueError(
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+ f"heads ({config.num_attention_heads})"
+ )
+
+ self.config = config
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+ self.attention_dropout = config.attention_probs_dropout_prob
+ self.scaling = self.attention_head_size**-0.5
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.attention_head_size)
+
+ query_states = self.query(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.key(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_states = self.value(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfOutput
+class ClapTextSelfOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+# Copied from transformers.models.align.modeling_align.AlignTextAttention with Align->Clap
+class ClapTextAttention(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.self = ClapTextSelfAttention(config)
+ self.output = ClapTextSelfOutput(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states, _ = self.self(
+ hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = self.output(hidden_states, residual)
+ return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate
+class ClapTextIntermediate(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertOutput
+class ClapTextOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+# Copied from transformers.models.align.modeling_align.AlignTextLayer with Align->Clap
+class ClapTextLayer(GradientCheckpointingLayer):
+ def __init__(self, config):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+ self.attention = ClapTextAttention(config)
+ self.intermediate = ClapTextIntermediate(config)
+ self.output = ClapTextOutput(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ hidden_states = self.attention(
+ hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+
+ hidden_states = apply_chunking_to_forward(
+ self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, hidden_states
+ )
+
+ return hidden_states
+
+ def feed_forward_chunk(self, attention_output):
+ intermediate_output = self.intermediate(attention_output)
+ layer_output = self.output(intermediate_output, attention_output)
+ return layer_output
+
+
+# Copied from transformers.models.align.modeling_align.AlignTextEncoder with Align->Clap
+class ClapTextEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList([ClapTextLayer(config) for i in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutput:
+ for layer_module in self.layer:
+ hidden_states = layer_module(
+ hidden_states,
+ attention_mask,
+ **kwargs,
+ )
+
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+# Copied from transformers.models.bert.modeling_bert.BertPooler
+class ClapTextPooler(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.activation = nn.Tanh()
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ # We "pool" the model by simply taking the hidden state corresponding
+ # to the first token.
+ first_token_tensor = hidden_states[:, 0]
+ pooled_output = self.dense(first_token_tensor)
+ pooled_output = self.activation(pooled_output)
+ return pooled_output
+
+
+@auto_docstring
+class ClapPreTrainedModel(PreTrainedModel):
+ config: ClapConfig
+ base_model_prefix = "clap"
+ input_modalities = ("audio", "text")
+ supports_gradient_checkpointing = False
+
+ @torch.no_grad()
+ def _init_weights(self, module: nn.Module):
+ """Initialize the weights"""
+ factor = self.config.initializer_factor
+
+ if isinstance(module, ClapTextEmbeddings):
+ init.normal_(module.position_embeddings.weight, mean=0.0, std=factor * 0.02)
+ init.normal_(module.token_type_embeddings.weight, mean=0.0, std=factor * 0.02)
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+ init.zeros_(module.token_type_ids)
+ elif isinstance(module, ClapModel):
+ init.constant_(module.logit_scale_a, math.log(self.config.logit_scale_init_value))
+ init.constant_(module.logit_scale_t, math.log(self.config.logit_scale_init_value))
+ elif isinstance(module, nn.Embedding):
+ init.normal_(module.weight, mean=0.0, std=factor * 0.02)
+ elif isinstance(module, (nn.LayerNorm, nn.BatchNorm2d)):
+ init.zeros_(module.bias)
+ init.ones_(module.weight)
+ if getattr(module, "running_mean", None) is not None:
+ init.zeros_(module.running_mean)
+ init.ones_(module.running_var)
+ init.zeros_(module.num_batches_tracked)
+ elif isinstance(module, (nn.Conv2d, nn.Linear)):
+ in_proj_std = (self.config.hidden_size**-0.5) * ((2 * self.config.num_hidden_layers) ** -0.5) * factor
+ init.normal_(module.weight, std=in_proj_std)
+ if module.bias is not None:
+ init.zeros_(module.bias)
+ elif isinstance(module, ClapAudioSelfAttention):
+ init.zeros_(module.relative_position_bias_table)
+ init.copy_(module.relative_position_index, module.create_relative_position_index())
+
+
+class ClapAudioModel(ClapPreTrainedModel):
+ config: ClapAudioConfig
+ main_input_name = "input_features"
+ input_modalities = "audio"
+
+ def __init__(self, config: ClapAudioConfig):
+ super().__init__(config)
+ self.audio_encoder = ClapAudioEncoder(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.audio_encoder.patch_embed.proj
+
+ @auto_docstring
+ def forward(
+ self,
+ input_features: torch.FloatTensor | None = None,
+ is_longer: torch.BoolTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ is_longer (`torch.FloatTensor`, of shape `(batch_size, 1)`, *optional*):
+ Whether the audio clip is longer than `max_length`. If `True`, a feature fusion will be enabled to enhance
+ the features.
+
+ Examples:
+
+ ```python
+ >>> from datasets import load_dataset
+ >>> from transformers import AutoProcessor, ClapAudioModel
+
+ >>> dataset = load_dataset("hf-internal-testing/ashraq-esc50-1-dog-example")
+ >>> audio_sample = dataset["train"]["audio"][0]["array"]
+
+ >>> model = ClapAudioModel.from_pretrained("laion/clap-htsat-fused")
+ >>> processor = AutoProcessor.from_pretrained("laion/clap-htsat-fused")
+
+ >>> inputs = processor(audio=audio_sample, return_tensors="pt")
+
+ >>> outputs = model(**inputs)
+ >>> last_hidden_state = outputs.last_hidden_state
+ ```"""
+ return self.audio_encoder(
+ input_features=input_features,
+ is_longer=is_longer,
+ **kwargs,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+ cross-attention is added between the self-attention layers, following the architecture described in *Attention is
+ all you need*_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
+ Kaiser and Illia Polosukhin.
+
+ To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set
+ to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and
+ `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.
+
+ .. _*Attention is all you need*: https://huggingface.co/papers/1706.03762
+ """
+)
+class ClapTextModel(ClapPreTrainedModel):
+ config: ClapTextConfig
+ input_modalities = ("text",)
+ _can_record_outputs = {
+ "hidden_states": ClapTextLayer,
+ "attentions": ClapTextSelfAttention,
+ }
+
+ def __init__(self, config, add_pooling_layer=True):
+ r"""
+ add_pooling_layer (bool, *optional*, defaults to `True`):
+ Whether to add a pooling layer
+ """
+ super().__init__(config)
+ self.config = config
+
+ self.embeddings = ClapTextEmbeddings(config)
+ self.encoder = ClapTextEncoder(config)
+
+ self.pooler = ClapTextPooler(config) if add_pooling_layer else None
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embeddings.word_embeddings
+
+ def set_input_embeddings(self, value):
+ self.embeddings.word_embeddings = value
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPooling:
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+ elif input_ids is not None:
+ self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
+ input_shape = input_ids.size()
+ elif inputs_embeds is not None:
+ input_shape = inputs_embeds.size()[:-1]
+ else:
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+ batch_size, seq_length = input_shape
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+ if attention_mask is None:
+ attention_mask = torch.ones(((batch_size, seq_length)), device=device)
+
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+ # ourselves in which case we just need to make it broadcastable to all heads.
+ extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)
+
+ embedding_output = self.embeddings(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ )
+ encoder_outputs = self.encoder(
+ embedding_output,
+ attention_mask=extended_attention_mask,
+ **kwargs,
+ )
+ sequence_output = encoder_outputs[0]
+ pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ )
+
+
+@auto_docstring
+class ClapModel(ClapPreTrainedModel):
+ config: ClapConfig
+
+ def __init__(self, config: ClapConfig):
+ super().__init__(config)
+
+ if not isinstance(config.text_config, ClapTextConfig):
+ raise TypeError(
+ "config.text_config is expected to be of type ClapTextConfig but is of type"
+ f" {type(config.text_config)}."
+ )
+
+ if not isinstance(config.audio_config, ClapAudioConfig):
+ raise TypeError(
+ "config.audio_config is expected to be of type ClapAudioConfig but is of type"
+ f" {type(config.audio_config)}."
+ )
+
+ text_config = config.text_config
+ audio_config = config.audio_config
+
+ self.logit_scale_a = nn.Parameter(torch.tensor(math.log(config.logit_scale_init_value)))
+ self.logit_scale_t = nn.Parameter(torch.tensor(math.log(config.logit_scale_init_value)))
+
+ self.projection_dim = config.projection_dim
+
+ self.text_model = ClapTextModel(text_config)
+ self.text_projection = ClapProjectionLayer(text_config)
+
+ self.audio_model = ClapAudioModel(audio_config)
+ self.audio_projection = ClapProjectionLayer(audio_config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def get_text_features(
+ self,
+ input_ids: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ Examples:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoTokenizer, ClapModel
+
+ >>> model = ClapModel.from_pretrained("laion/clap-htsat-unfused")
+ >>> tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused")
+
+ >>> inputs = tokenizer(["the sound of a cat", "the sound of a dog"], padding=True, return_tensors="pt")
+ >>> with torch.inference_mode():
+ ... text_features = model.get_text_features(**inputs)
+ ```"""
+ text_outputs: BaseModelOutputWithPooling = self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ **kwargs,
+ )
+ text_features = self.text_projection(text_outputs.pooler_output)
+ text_outputs.pooler_output = F.normalize(text_features, dim=-1)
+
+ return text_outputs
+
+ @can_return_tuple
+ @auto_docstring
+ def get_audio_features(
+ self,
+ input_features: torch.Tensor,
+ is_longer: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ is_longer (`torch.FloatTensor`, of shape `(batch_size, 1)`, *optional*):
+ Whether the audio clip is longer than `max_length`. If `True`, a feature fusion will be enabled to enhance
+ the features.
+
+ Examples:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoFeatureExtractor, ClapModel
+
+ >>> model = ClapModel.from_pretrained("laion/clap-htsat-unfused")
+ >>> feature_extractor = AutoFeatureExtractor.from_pretrained("laion/clap-htsat-unfused")
+ >>> random_audio = torch.rand((16_000))
+
+ >>> inputs = feature_extractor(random_audio, return_tensors="pt")
+ >>> with torch.inference_mode():
+ ... audio_features = model.get_audio_features(**inputs)
+ ```"""
+ audio_outputs: BaseModelOutputWithPooling = self.audio_model(
+ input_features=input_features, is_longer=is_longer, **kwargs
+ )
+ audio_features = self.audio_projection(audio_outputs.pooler_output)
+ audio_outputs.pooler_output = F.normalize(audio_features, dim=-1)
+
+ return audio_outputs
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ input_features: torch.FloatTensor | None = None,
+ is_longer: torch.BoolTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ return_loss: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | ClapOutput:
+ r"""
+ is_longer (`torch.FloatTensor`, of shape `(batch_size, 1)`, *optional*):
+ Whether the audio clip is longer than `max_length`. If `True`, a feature fusion will be enabled to enhance
+ the features.
+ return_loss (`bool`, *optional*):
+ Whether or not to return the contrastive loss.
+
+ Examples:
+
+ ```python
+ >>> from datasets import load_dataset
+ >>> from transformers import AutoProcessor, ClapModel
+
+ >>> dataset = load_dataset("hf-internal-testing/ashraq-esc50-1-dog-example")
+ >>> audio_sample = dataset["train"]["audio"][0]["array"]
+
+ >>> model = ClapModel.from_pretrained("laion/clap-htsat-unfused")
+ >>> processor = AutoProcessor.from_pretrained("laion/clap-htsat-unfused")
+
+ >>> input_text = ["Sound of a dog", "Sound of vacuum cleaner"]
+
+ >>> inputs = processor(text=input_text, audio=audio_sample, return_tensors="pt", padding=True)
+
+ >>> outputs = model(**inputs)
+ >>> logits_per_audio = outputs.logits_per_audio # this is the audio-text similarity score
+ >>> probs = logits_per_audio.softmax(dim=-1) # we can take the softmax to get the label probabilities
+ ```"""
+ audio_outputs = self.audio_model(
+ input_features=input_features,
+ is_longer=is_longer,
+ **kwargs,
+ )
+
+ text_outputs = self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+ audio_embeds = audio_outputs.pooler_output
+ audio_embeds = self.audio_projection(audio_embeds)
+
+ text_embeds = text_outputs.pooler_output
+ text_embeds = self.text_projection(text_embeds)
+
+ # normalized features
+ audio_embeds = audio_embeds / audio_embeds.norm(p=2, dim=-1, keepdim=True)
+ text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
+
+ # cosine similarity as logits
+ logit_scale_text = self.logit_scale_t.exp()
+ logit_scale_audio = self.logit_scale_a.exp()
+ logits_per_text = torch.matmul(text_embeds, audio_embeds.t()) * logit_scale_text
+ logits_per_audio = torch.matmul(audio_embeds, text_embeds.t()) * logit_scale_audio
+
+ loss = None
+ if return_loss:
+ caption_loss = contrastive_loss(logits_per_text)
+ audio_loss = contrastive_loss(logits_per_audio.t())
+ loss = (caption_loss + audio_loss) / 2.0
+
+ return ClapOutput(
+ loss=loss,
+ logits_per_audio=logits_per_audio,
+ logits_per_text=logits_per_text,
+ text_embeds=text_embeds,
+ audio_embeds=audio_embeds,
+ text_model_output=text_outputs,
+ audio_model_output=audio_outputs,
+ )
+
+
+@auto_docstring
+class ClapTextModelWithProjection(ClapPreTrainedModel):
+ config: ClapTextConfig
+ input_modalities = ("text",)
+ _can_record_outputs = {
+ "hidden_states": ClapTextLayer,
+ "attentions": ClapTextSelfAttention,
+ }
+
+ def __init__(self, config: ClapTextConfig):
+ super().__init__(config)
+ self.text_model = ClapTextModel(config)
+ self.text_projection = ClapProjectionLayer(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.text_model.embeddings.word_embeddings
+
+ def set_input_embeddings(self, value):
+ self.text_model.embeddings.word_embeddings = value
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | ClapTextModelOutput:
+ r"""
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, ClapTextModelWithProjection
+
+ >>> model = ClapTextModelWithProjection.from_pretrained("laion/clap-htsat-unfused")
+ >>> tokenizer = AutoTokenizer.from_pretrained("laion/clap-htsat-unfused")
+
+ >>> inputs = tokenizer(["a sound of a cat", "a sound of a dog"], padding=True, return_tensors="pt")
+
+ >>> outputs = model(**inputs)
+ >>> text_embeds = outputs.text_embeds
+ ```"""
+ text_outputs: BaseModelOutputWithPooling = self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ **kwargs,
+ )
+ pooled_output = text_outputs.pooler_output
+ text_embeds = self.text_projection(pooled_output)
+
+ return ClapTextModelOutput(
+ text_embeds=text_embeds,
+ last_hidden_state=text_outputs.last_hidden_state,
+ hidden_states=text_outputs.hidden_states,
+ attentions=text_outputs.attentions,
+ )
+
+
+@auto_docstring
+class ClapAudioModelWithProjection(ClapPreTrainedModel):
+ config: ClapAudioConfig
+ main_input_name = "input_features"
+ input_modalities = "audio"
+
+ def __init__(self, config: ClapAudioConfig):
+ super().__init__(config)
+ self.audio_model = ClapAudioModel(config)
+ self.audio_projection = ClapProjectionLayer(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.audio_model.audio_encoder.patch_embed.proj
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_features: torch.FloatTensor | None = None,
+ is_longer: torch.BoolTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | ClapAudioModelOutput:
+ r"""
+ is_longer (`torch.FloatTensor`, of shape `(batch_size, 1)`, *optional*):
+ Whether the audio clip is longer than `max_length`. If `True`, a feature fusion will be enabled to enhance
+ the features.
+
+ Examples:
+
+ ```python
+ >>> from datasets import load_dataset
+ >>> from transformers import ClapAudioModelWithProjection, ClapProcessor
+
+ >>> model = ClapAudioModelWithProjection.from_pretrained("laion/clap-htsat-fused")
+ >>> processor = ClapProcessor.from_pretrained("laion/clap-htsat-fused")
+
+ >>> dataset = load_dataset("hf-internal-testing/ashraq-esc50-1-dog-example")
+ >>> audio_sample = dataset["train"]["audio"][0]["array"]
+
+ >>> inputs = processor(audio=audio_sample, return_tensors="pt")
+ >>> outputs = model(**inputs)
+ >>> audio_embeds = outputs.audio_embeds
+ ```"""
+ audio_outputs: BaseModelOutputWithPooling = self.audio_model(
+ input_features=input_features,
+ is_longer=is_longer,
+ **kwargs,
+ )
+
+ audio_embeds = self.audio_projection(audio_outputs.pooler_output)
+
+ return ClapAudioModelOutput(
+ audio_embeds=audio_embeds,
+ last_hidden_state=audio_outputs.last_hidden_state,
+ attentions=audio_outputs.attentions,
+ hidden_states=audio_outputs.hidden_states,
+ )
+
+
+__all__ = [
+ "ClapModel",
+ "ClapPreTrainedModel",
+ "ClapTextModel",
+ "ClapTextModelWithProjection",
+ "ClapAudioModel",
+ "ClapAudioModelWithProjection",
+]
diff --git a/third_party/transformers/src/transformers/models/clap/processing_clap.py b/third_party/transformers/src/transformers/models/clap/processing_clap.py
new file mode 100644
index 0000000000000000000000000000000000000000..567aa19ce60775e9f14ed220f4edecb910cf1c21
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/clap/processing_clap.py
@@ -0,0 +1,31 @@
+# Copyright 2023 The HuggingFace Inc. team.
+#
+# 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.
+"""
+Audio/Text processor class for CLAP
+"""
+
+from ...processing_utils import ProcessorMixin
+from ...utils import auto_docstring, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring
+class ClapProcessor(ProcessorMixin):
+ def __init__(self, feature_extractor, tokenizer):
+ super().__init__(feature_extractor, tokenizer)
+
+
+__all__ = ["ClapProcessor"]
diff --git a/third_party/transformers/src/transformers/models/clipseg/__init__.py b/third_party/transformers/src/transformers/models/clipseg/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..55b38987fd0a3e443b8ae94ffbdd6f73c79ba619
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/clipseg/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_clipseg import *
+ from .modeling_clipseg import *
+ from .processing_clipseg import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/clipseg/configuration_clipseg.py b/third_party/transformers/src/transformers/models/clipseg/configuration_clipseg.py
new file mode 100644
index 0000000000000000000000000000000000000000..278671022f3b2d8e591da7738d98d9bd06ca686f
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/clipseg/configuration_clipseg.py
@@ -0,0 +1,238 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""CLIPSeg model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="CIDAS/clipseg-rd64")
+@strict
+class CLIPSegTextConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import CLIPSegTextConfig, CLIPSegTextModel
+
+ >>> # Initializing a CLIPSegTextConfig with CIDAS/clipseg-rd64 style configuration
+ >>> configuration = CLIPSegTextConfig()
+
+ >>> # Initializing a CLIPSegTextModel (with random weights) from the CIDAS/clipseg-rd64 style configuration
+ >>> model = CLIPSegTextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "clipseg_text_model"
+ base_config_key = "text_config"
+
+ vocab_size: int = 49408
+ hidden_size: int = 512
+ intermediate_size: int = 2048
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 8
+ max_position_embeddings: int = 77
+ hidden_act: str = "quick_gelu"
+ layer_norm_eps: float = 1e-5
+ attention_dropout: float | int = 0.0
+ initializer_range: float = 0.02
+ initializer_factor: float = 1.0
+ pad_token_id: int | None = 1
+ bos_token_id: int | None = 49406
+ eos_token_id: int | list[int] | None = 49407
+
+
+@auto_docstring(checkpoint="CIDAS/clipseg-rd64")
+@strict
+class CLIPSegVisionConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import CLIPSegVisionConfig, CLIPSegVisionModel
+
+ >>> # Initializing a CLIPSegVisionConfig with CIDAS/clipseg-rd64 style configuration
+ >>> configuration = CLIPSegVisionConfig()
+
+ >>> # Initializing a CLIPSegVisionModel (with random weights) from the CIDAS/clipseg-rd64 style configuration
+ >>> model = CLIPSegVisionModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "clipseg_vision_model"
+ base_config_key = "vision_config"
+
+ hidden_size: int = 768
+ intermediate_size: int = 3072
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 12
+ num_channels: int = 3
+ image_size: int | list[int] | tuple[int, int] = 224
+ patch_size: int | list[int] | tuple[int, int] = 32
+ hidden_act: str = "quick_gelu"
+ layer_norm_eps: float = 1e-5
+ attention_dropout: float | int = 0.0
+ initializer_range: float = 0.02
+ initializer_factor: float = 1.0
+
+
+@auto_docstring(checkpoint="CIDAS/clipseg-rd64")
+@strict
+class CLIPSegConfig(PreTrainedConfig):
+ r"""
+ extract_layers (`list[int]`, *optional*, defaults to `[3, 6, 9]`):
+ Layers to extract when forwarding the query image through the frozen visual backbone of CLIP.
+ reduce_dim (`int`, *optional*, defaults to 64):
+ Dimensionality to reduce the CLIP vision embedding.
+ conditional_layer (`int`, *optional*, defaults to 0):
+ The layer to use of the Transformer encoder whose activations will be combined with the condition
+ embeddings using FiLM (Feature-wise Linear Modulation). If 0, the last layer is used.
+ use_complex_transposed_convolution (`bool`, *optional*, defaults to `False`):
+ Whether to use a more complex transposed convolution in the decoder, enabling more fine-grained
+ segmentation..
+
+ Example:
+
+ ```python
+ >>> from transformers import CLIPSegConfig, CLIPSegModel
+
+ >>> # Initializing a CLIPSegConfig with CIDAS/clipseg-rd64 style configuration
+ >>> configuration = CLIPSegConfig()
+
+ >>> # Initializing a CLIPSegModel (with random weights) from the CIDAS/clipseg-rd64 style configuration
+ >>> model = CLIPSegModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+
+ >>> # We can also initialize a CLIPSegConfig from a CLIPSegTextConfig and a CLIPSegVisionConfig
+
+ >>> # Initializing a CLIPSegText and CLIPSegVision configuration
+ >>> config_text = CLIPSegTextConfig()
+ >>> config_vision = CLIPSegVisionConfig()
+
+ >>> config = CLIPSegConfig(text_config=config_text, vision_config=config_vision)
+ ```"""
+
+ model_type = "clipseg"
+ sub_configs = {"text_config": CLIPSegTextConfig, "vision_config": CLIPSegVisionConfig}
+
+ text_config: dict | CLIPSegTextConfig | None = None
+ vision_config: dict | CLIPSegVisionConfig | None = None
+ projection_dim: int | None = 512
+ logit_scale_init_value: float | int | None = 2.6592
+ initializer_factor: float | None = 1.0
+ extract_layers: list[int] | tuple[int, ...] = (3, 6, 9)
+ reduce_dim: int = 64
+ decoder_num_attention_heads: int = 4
+ decoder_attention_dropout: float | int = 0.0
+ decoder_hidden_act: str = "quick_gelu"
+ decoder_intermediate_size: int = 2048
+ conditional_layer: int = 0
+ use_complex_transposed_convolution: bool = False
+
+ def __post_init__(self, **kwargs):
+ if self.text_config is None:
+ text_config = {}
+ logger.info("`text_config` is `None`. Initializing the `CLIPSegTextConfig` with default values.")
+ elif isinstance(self.text_config, CLIPSegTextConfig):
+ text_config = self.text_config.to_dict()
+ else:
+ text_config = self.text_config
+
+ if self.vision_config is None:
+ vision_config = {}
+ logger.info("`vision_config` is `None`. initializing the `CLIPSegVisionConfig` with default values.")
+ elif isinstance(self.vision_config, CLIPSegVisionConfig):
+ vision_config = self.vision_config.to_dict()
+ else:
+ vision_config = self.vision_config
+
+ # For backward compatibility check keyword args
+ # Instead of simply assigning `[text|vision]_config_dict` to `[text|vision]_config`, we use the values in
+ # `[text|vision]_config_dict` to update the values in `[text|vision]_config`. The values should be same in most
+ # cases, but we don't want to break anything regarding `_config_dict` that existed before commit `8827e1b2`.
+ text_config_dict = kwargs.pop("text_config_dict", None)
+ vision_config_dict = kwargs.pop("vision_config_dict", None)
+
+ if text_config_dict is not None:
+ # This is the complete result when using `text_config_dict`.
+ _text_config_dict = CLIPSegTextConfig(**text_config_dict).to_dict()
+
+ # Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
+ for key, value in _text_config_dict.items():
+ if key in text_config and value != text_config[key] and key != "transformers_version":
+ # If specified in `text_config_dict`
+ if key in text_config_dict:
+ message = (
+ f"`{key}` is found in both `text_config_dict` and `text_config` but with different values. "
+ f'The value `text_config_dict["{key}"]` will be used instead.'
+ )
+ # If inferred from default argument values (just to be super careful)
+ else:
+ message = (
+ f"`text_config_dict` is provided which will be used to initialize `CLIPSegTextConfig`. The "
+ f'value `text_config["{key}"]` will be overridden.'
+ )
+ logger.info(message)
+
+ # Update all values in `text_config` with the ones in `_text_config_dict`.
+ text_config.update(_text_config_dict)
+
+ if vision_config_dict is not None:
+ # This is the complete result when using `vision_config_dict`.
+ _vision_config_dict = CLIPSegVisionConfig(**vision_config_dict).to_dict()
+ # convert keys to string instead of integer
+ if "id2label" in _vision_config_dict:
+ _vision_config_dict["id2label"] = {
+ str(key): value for key, value in _vision_config_dict["id2label"].items()
+ }
+
+ # Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
+ for key, value in _vision_config_dict.items():
+ if key in vision_config and value != vision_config[key] and key != "transformers_version":
+ # If specified in `vision_config_dict`
+ if key in vision_config_dict:
+ message = (
+ f"`{key}` is found in both `vision_config_dict` and `vision_config` but with different "
+ f'values. The value `vision_config_dict["{key}"]` will be used instead.'
+ )
+ # If inferred from default argument values (just to be super careful)
+ else:
+ message = (
+ f"`vision_config_dict` is provided which will be used to initialize `CLIPSegVisionConfig`. "
+ f'The value `vision_config["{key}"]` will be overridden.'
+ )
+ logger.info(message)
+
+ # Update all values in `vision_config` with the ones in `_vision_config_dict`.
+ vision_config.update(_vision_config_dict)
+
+ # Finally we can convert back our unified text/vision configs to `PretrainedConfig`
+ self.text_config = CLIPSegTextConfig(**text_config)
+ self.vision_config = CLIPSegVisionConfig(**vision_config)
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["CLIPSegConfig", "CLIPSegTextConfig", "CLIPSegVisionConfig"]
diff --git a/third_party/transformers/src/transformers/models/clipseg/convert_clipseg_original_pytorch_to_hf.py b/third_party/transformers/src/transformers/models/clipseg/convert_clipseg_original_pytorch_to_hf.py
new file mode 100644
index 0000000000000000000000000000000000000000..ca3513ac4a5a81fd92c316b00b47da3d0eb0e75e
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/clipseg/convert_clipseg_original_pytorch_to_hf.py
@@ -0,0 +1,267 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+
+"""Convert CLIPSeg checkpoints from the original repository. URL: https://github.com/timojl/clipseg."""
+
+import argparse
+from io import BytesIO
+
+import httpx
+import torch
+from PIL import Image
+
+from transformers import (
+ CLIPSegConfig,
+ CLIPSegForImageSegmentation,
+ CLIPSegProcessor,
+ CLIPSegTextConfig,
+ CLIPSegVisionConfig,
+ CLIPTokenizer,
+ ViTImageProcessor,
+)
+
+
+def get_clipseg_config(model_name):
+ text_config = CLIPSegTextConfig()
+ vision_config = CLIPSegVisionConfig(patch_size=16)
+
+ use_complex_transposed_convolution = "refined" in model_name
+ reduce_dim = 16 if "rd16" in model_name else 64
+
+ config = CLIPSegConfig(
+ text_config=text_config,
+ vision_config=vision_config,
+ use_complex_transposed_convolution=use_complex_transposed_convolution,
+ reduce_dim=reduce_dim,
+ )
+ return config
+
+
+def rename_key(name):
+ # update prefixes
+ if "clip_model" in name:
+ name = name.replace("clip_model", "clip")
+ if "transformer" in name:
+ if "visual" in name:
+ name = name.replace("visual.transformer", "vision_model")
+ else:
+ name = name.replace("transformer", "text_model")
+ if "resblocks" in name:
+ name = name.replace("resblocks", "encoder.layers")
+ if "ln_1" in name:
+ name = name.replace("ln_1", "layer_norm1")
+ if "ln_2" in name:
+ name = name.replace("ln_2", "layer_norm2")
+ if "c_fc" in name:
+ name = name.replace("c_fc", "fc1")
+ if "c_proj" in name:
+ name = name.replace("c_proj", "fc2")
+ if "attn" in name and "self" not in name:
+ name = name.replace("attn", "self_attn")
+ # text encoder
+ if "token_embedding" in name:
+ name = name.replace("token_embedding", "text_model.embeddings.token_embedding")
+ if "positional_embedding" in name and "visual" not in name:
+ name = name.replace("positional_embedding", "text_model.embeddings.position_embedding.weight")
+ if "ln_final" in name:
+ name = name.replace("ln_final", "text_model.final_layer_norm")
+ # vision encoder
+ if "visual.class_embedding" in name:
+ name = name.replace("visual.class_embedding", "vision_model.embeddings.class_embedding")
+ if "visual.conv1" in name:
+ name = name.replace("visual.conv1", "vision_model.embeddings.patch_embedding")
+ if "visual.positional_embedding" in name:
+ name = name.replace("visual.positional_embedding", "vision_model.embeddings.position_embedding.weight")
+ if "visual.ln_pre" in name:
+ name = name.replace("visual.ln_pre", "vision_model.pre_layrnorm")
+ if "visual.ln_post" in name:
+ name = name.replace("visual.ln_post", "vision_model.post_layernorm")
+ # projection layers
+ if "visual.proj" in name:
+ name = name.replace("visual.proj", "visual_projection.weight")
+ if "text_projection" in name:
+ name = name.replace("text_projection", "text_projection.weight")
+ # decoder
+ if "trans_conv" in name:
+ name = name.replace("trans_conv", "transposed_convolution")
+ if "film_mul" in name or "film_add" in name or "reduce" in name or "transposed_convolution" in name:
+ name = "decoder." + name
+ if "blocks" in name:
+ name = name.replace("blocks", "decoder.layers")
+ if "linear1" in name:
+ name = name.replace("linear1", "mlp.fc1")
+ if "linear2" in name:
+ name = name.replace("linear2", "mlp.fc2")
+ if "norm1" in name and "layer_" not in name:
+ name = name.replace("norm1", "layer_norm1")
+ if "norm2" in name and "layer_" not in name:
+ name = name.replace("norm2", "layer_norm2")
+
+ return name
+
+
+def convert_state_dict(orig_state_dict, config):
+ for key in orig_state_dict.copy():
+ val = orig_state_dict.pop(key)
+
+ if key.startswith("clip_model") and "attn.in_proj" in key:
+ key_split = key.split(".")
+ if "visual" in key:
+ layer_num = int(key_split[4])
+ dim = config.vision_config.hidden_size
+ prefix = "vision_model"
+ else:
+ layer_num = int(key_split[3])
+ dim = config.text_config.hidden_size
+ prefix = "text_model"
+
+ if "weight" in key:
+ orig_state_dict[f"clip.{prefix}.encoder.layers.{layer_num}.self_attn.q_proj.weight"] = val[:dim, :]
+ orig_state_dict[f"clip.{prefix}.encoder.layers.{layer_num}.self_attn.k_proj.weight"] = val[
+ dim : dim * 2, :
+ ]
+ orig_state_dict[f"clip.{prefix}.encoder.layers.{layer_num}.self_attn.v_proj.weight"] = val[-dim:, :]
+ else:
+ orig_state_dict[f"clip.{prefix}.encoder.layers.{layer_num}.self_attn.q_proj.bias"] = val[:dim]
+ orig_state_dict[f"clip.{prefix}.encoder.layers.{layer_num}.self_attn.k_proj.bias"] = val[dim : dim * 2]
+ orig_state_dict[f"clip.{prefix}.encoder.layers.{layer_num}.self_attn.v_proj.bias"] = val[-dim:]
+ elif "self_attn" in key and "out_proj" not in key:
+ key_split = key.split(".")
+ layer_num = int(key_split[1])
+ dim = config.reduce_dim
+ if "weight" in key:
+ orig_state_dict[f"decoder.layers.{layer_num}.self_attn.q_proj.weight"] = val[:dim, :]
+ orig_state_dict[f"decoder.layers.{layer_num}.self_attn.k_proj.weight"] = val[dim : dim * 2, :]
+ orig_state_dict[f"decoder.layers.{layer_num}.self_attn.v_proj.weight"] = val[-dim:, :]
+ else:
+ orig_state_dict[f"decoder.layers.{layer_num}.self_attn.q_proj.bias"] = val[:dim]
+ orig_state_dict[f"decoder.layers.{layer_num}.self_attn.k_proj.bias"] = val[dim : dim * 2]
+ orig_state_dict[f"decoder.layers.{layer_num}.self_attn.v_proj.bias"] = val[-dim:]
+ else:
+ new_name = rename_key(key)
+ if "visual_projection" in new_name or "text_projection" in new_name:
+ val = val.T
+ orig_state_dict[new_name] = val
+
+ return orig_state_dict
+
+
+# We will verify our results on an image of cute cats
+def prepare_img():
+ url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ with httpx.stream("GET", url) as response:
+ image = Image.open(BytesIO(response.read()))
+ return image
+
+
+def convert_clipseg_checkpoint(model_name, checkpoint_path, pytorch_dump_folder_path, push_to_hub):
+ config = get_clipseg_config(model_name)
+ model = CLIPSegForImageSegmentation(config)
+ model.eval()
+
+ state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+
+ # remove some keys
+ for key in state_dict.copy():
+ if key.startswith("model"):
+ state_dict.pop(key, None)
+
+ # rename some keys
+ state_dict = convert_state_dict(state_dict, config)
+ missing_keys, unexpected_keys = model.load_state_dict(state_dict, strict=False)
+
+ if missing_keys != ["clip.text_model.embeddings.position_ids", "clip.vision_model.embeddings.position_ids"]:
+ raise ValueError(f"Missing keys that are not expected: {missing_keys}")
+ if unexpected_keys != ["decoder.reduce.weight", "decoder.reduce.bias"]:
+ raise ValueError(f"Unexpected keys: {unexpected_keys}")
+
+ image_processor = ViTImageProcessor(size=352)
+ tokenizer = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
+ processor = CLIPSegProcessor(image_processor=image_processor, tokenizer=tokenizer)
+
+ image = prepare_img()
+ text = ["a glass", "something to fill", "wood", "a jar"]
+
+ inputs = processor(text=text, images=[image] * len(text), padding="max_length", return_tensors="pt")
+
+ with torch.no_grad():
+ outputs = model(**inputs)
+
+ # verify values
+ expected_conditional = torch.tensor([0.1110, -0.1882, 0.1645])
+ expected_pooled_output = torch.tensor([0.2692, -0.7197, -0.1328])
+ if model_name == "clipseg-rd64-refined":
+ expected_masks_slice = torch.tensor(
+ [[-10.0407, -9.9431, -10.2646], [-9.9751, -9.7064, -9.9586], [-9.6891, -9.5645, -9.9618]]
+ )
+ elif model_name == "clipseg-rd64":
+ expected_masks_slice = torch.tensor(
+ [[-7.2877, -7.2711, -7.2463], [-7.2652, -7.2780, -7.2520], [-7.2239, -7.2204, -7.2001]]
+ )
+ elif model_name == "clipseg-rd16":
+ expected_masks_slice = torch.tensor(
+ [[-6.3955, -6.4055, -6.4151], [-6.3911, -6.4033, -6.4100], [-6.3474, -6.3702, -6.3762]]
+ )
+ else:
+ raise ValueError(f"Model name {model_name} not supported.")
+
+ assert torch.allclose(outputs.logits[0, :3, :3], expected_masks_slice, atol=1e-3)
+ assert torch.allclose(outputs.conditional_embeddings[0, :3], expected_conditional, atol=1e-3)
+ assert torch.allclose(outputs.pooled_output[0, :3], expected_pooled_output, atol=1e-3)
+ print("Looks ok!")
+
+ if pytorch_dump_folder_path is not None:
+ print(f"Saving model and processor to {pytorch_dump_folder_path}")
+ model.save_pretrained(pytorch_dump_folder_path)
+ processor.save_pretrained(pytorch_dump_folder_path)
+
+ if push_to_hub:
+ print(f"Pushing model and processor for {model_name} to the hub")
+ model.push_to_hub(f"CIDAS/{model_name}")
+ processor.push_to_hub(f"CIDAS/{model_name}")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ # Required parameters
+ parser.add_argument(
+ "--model_name",
+ default="clipseg-rd64",
+ type=str,
+ choices=["clipseg-rd16", "clipseg-rd64", "clipseg-rd64-refined"],
+ help=(
+ "Name of the model. Supported models are: clipseg-rd64, clipseg-rd16 and clipseg-rd64-refined (rd meaning"
+ " reduce dimension)"
+ ),
+ )
+ parser.add_argument(
+ "--checkpoint_path",
+ default="/Users/nielsrogge/Documents/CLIPSeg/clip_plus_rd64-uni.pth",
+ type=str,
+ help=(
+ "Path to the original checkpoint. Note that the script assumes that the checkpoint includes both CLIP and"
+ " the decoder weights."
+ ),
+ )
+ parser.add_argument(
+ "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
+ )
+ parser.add_argument(
+ "--push_to_hub",
+ action="store_true",
+ help="Whether or not to push the converted model to the Hugging Face hub.",
+ )
+
+ args = parser.parse_args()
+ convert_clipseg_checkpoint(args.model_name, args.checkpoint_path, args.pytorch_dump_folder_path, args.push_to_hub)
diff --git a/third_party/transformers/src/transformers/models/clipseg/modeling_clipseg.py b/third_party/transformers/src/transformers/models/clipseg/modeling_clipseg.py
new file mode 100644
index 0000000000000000000000000000000000000000..ffee94dd5b95d47bffbc575a74037ad4284295c4
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/clipseg/modeling_clipseg.py
@@ -0,0 +1,1241 @@
+# Copyright 2022 The OpenAI Team Authors and The HuggingFace Team. All rights reserved.
+#
+# 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 CLIPSeg model."""
+
+import copy
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...masking_utils import create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, logging, torch_int
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_clipseg import CLIPSegConfig, CLIPSegTextConfig, CLIPSegVisionConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# contrastive loss function, adapted from
+# https://sachinruk.github.io/blog/pytorch/pytorch%20lightning/loss%20function/gpu/2021/03/07/CLIP.html
+def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:
+ return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device))
+
+
+# Copied from transformers.models.clip.modeling_clip.clip_loss with clip->clipseg
+def clipseg_loss(similarity: torch.Tensor) -> torch.Tensor:
+ caption_loss = contrastive_loss(similarity)
+ image_loss = contrastive_loss(similarity.t())
+ return (caption_loss + image_loss) / 2.0
+
+
+@dataclass
+@auto_docstring
+# Copied from transformers.models.clip.modeling_clip.CLIPOutput with CLIP->CLIPSeg
+class CLIPSegOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`):
+ Contrastive loss for image-text similarity.
+ logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`):
+ The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text
+ similarity scores.
+ logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`):
+ The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image
+ similarity scores.
+ text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):
+ The text embeddings obtained by applying the projection layer to the pooled output of [`CLIPSegTextModel`].
+ image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):
+ The image embeddings obtained by applying the projection layer to the pooled output of [`CLIPSegVisionModel`].
+ text_model_output (`BaseModelOutputWithPooling`):
+ The output of the [`CLIPSegTextModel`].
+ vision_model_output (`BaseModelOutputWithPooling`):
+ The output of the [`CLIPSegVisionModel`].
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits_per_image: torch.FloatTensor | None = None
+ logits_per_text: torch.FloatTensor | None = None
+ text_embeds: torch.FloatTensor | None = None
+ image_embeds: torch.FloatTensor | None = None
+ text_model_output: BaseModelOutputWithPooling = None
+ vision_model_output: BaseModelOutputWithPooling = None
+
+ def to_tuple(self) -> tuple[Any]:
+ return tuple(v.to_tuple() if isinstance(v, ModelOutput) else v for v in self.values())
+
+
+@dataclass
+@auto_docstring
+class CLIPSegDecoderOutput(ModelOutput):
+ r"""
+ logits (`torch.FloatTensor` of shape `(batch_size, height, width)`):
+ Classification scores for each pixel.
+ """
+
+ logits: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+
+
+@dataclass
+@auto_docstring
+class CLIPSegImageSegmentationOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Binary cross entropy loss for segmentation.
+ logits (`torch.FloatTensor` of shape `(batch_size, height, width)`):
+ Classification scores for each pixel.
+ conditional_embeddings (`torch.FloatTensor` of shape `(batch_size, projection_dim)`):
+ Conditional embeddings used for segmentation.
+ pooled_output (`torch.FloatTensor` of shape `(batch_size, embed_dim)`):
+ Pooled output of the [`CLIPSegVisionModel`].
+ vision_model_output (`BaseModelOutputWithPooling`):
+ The output of the [`CLIPSegVisionModel`].
+ decoder_output (`CLIPSegDecoderOutput`):
+ The output of the [`CLIPSegDecoder`].
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ conditional_embeddings: torch.FloatTensor | None = None
+ pooled_output: torch.FloatTensor | None = None
+ vision_model_output: BaseModelOutputWithPooling = None
+ decoder_output: CLIPSegDecoderOutput = None
+
+ def to_tuple(self) -> tuple[Any]:
+ return tuple(v.to_tuple() if isinstance(v, ModelOutput) else v for v in self.values())
+
+
+class CLIPSegVisionEmbeddings(nn.Module):
+ # Copied from transformers.models.clip.modeling_clip.CLIPVisionEmbeddings.__init__ with CLIP->CLIPSeg
+ def __init__(self, config: CLIPSegVisionConfig):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.image_size = config.image_size
+ self.patch_size = config.patch_size
+
+ self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))
+
+ self.patch_embedding = nn.Conv2d(
+ in_channels=config.num_channels,
+ out_channels=self.embed_dim,
+ kernel_size=self.patch_size,
+ stride=self.patch_size,
+ bias=False,
+ )
+
+ self.num_patches = (self.image_size // self.patch_size) ** 2
+ self.num_positions = self.num_patches + 1
+ self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+ self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)
+
+ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+ """
+ This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution
+ images. This method is also adapted to support torch.jit tracing.
+
+ Adapted from:
+ - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
+ - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
+ """
+
+ num_patches = embeddings.shape[1] - 1
+ position_embedding = self.position_embedding.weight.unsqueeze(0)
+ num_positions = position_embedding.shape[1] - 1
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embedding(self.position_ids)
+
+ class_pos_embed = position_embedding[:, :1]
+ patch_pos_embed = position_embedding[:, 1:]
+
+ dim = embeddings.shape[-1]
+
+ new_height = height // self.patch_size
+ new_width = width // self.patch_size
+
+ sqrt_num_positions = torch_int(num_positions**0.5)
+ patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
+ patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed,
+ size=(new_height, new_width),
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+ return torch.cat((class_pos_embed, patch_pos_embed), dim=1)
+
+ def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding=True) -> torch.Tensor:
+ batch_size, _, height, width = pixel_values.shape
+ if not interpolate_pos_encoding and (height != self.image_size or width != self.image_size):
+ raise ValueError(
+ f"Input image size ({height}*{width}) doesn't match model ({self.image_size}*{self.image_size})."
+ )
+ patch_embeds = self.patch_embedding(pixel_values) # shape = [*, width, grid, grid]
+ patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+
+ class_embeds = self.class_embedding.expand(batch_size, 1, -1)
+ embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
+ if interpolate_pos_encoding:
+ embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+ else:
+ embeddings = embeddings + self.position_embedding(self.position_ids)
+ return embeddings
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPTextEmbeddings with CLIP->CLIPSeg
+class CLIPSegTextEmbeddings(nn.Module):
+ def __init__(self, config: CLIPSegTextConfig):
+ super().__init__()
+ embed_dim = config.hidden_size
+
+ self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)
+ self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim)
+
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+ self.register_buffer(
+ "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+ )
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ ) -> torch.Tensor:
+ seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]
+ max_position_embedding = self.position_embedding.weight.shape[0]
+
+ if seq_length > max_position_embedding:
+ raise ValueError(
+ f"Sequence length must be less than max_position_embeddings (got `sequence length`: "
+ f"{seq_length} and max_position_embeddings: {max_position_embedding}"
+ )
+
+ if position_ids is None:
+ position_ids = self.position_ids[:, :seq_length]
+
+ if inputs_embeds is None:
+ inputs_embeds = self.token_embedding(input_ids)
+
+ position_embeddings = self.position_embedding(position_ids)
+ embeddings = inputs_embeds + position_embeddings
+
+ return embeddings
+
+
+# Copied from transformers.models.siglip.modeling_siglip.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class CLIPSegAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: CLIPSegVisionConfig | CLIPSegTextConfig):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = self.embed_dim // self.num_heads
+ if self.head_dim * self.num_heads != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+ f" {self.num_heads})."
+ )
+ self.scale = self.head_dim**-0.5
+ self.dropout = config.attention_dropout
+ self.is_causal = False
+
+ self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Input shape: Batch x Time x Channel"""
+ batch_size, seq_length, embed_dim = hidden_states.shape
+
+ queries = self.q_proj(hidden_states)
+ keys = self.k_proj(hidden_states)
+ values = self.v_proj(hidden_states)
+
+ queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+ keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+ values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ queries,
+ keys,
+ values,
+ attention_mask,
+ scaling=self.scale,
+ dropout=0.0 if not self.training else self.dropout,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->CLIPSeg
+class CLIPSegMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.activation_fn = ACT2FN[config.hidden_act]
+ self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+ self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->CLIPSeg
+class CLIPSegEncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: CLIPSegConfig):
+ super().__init__()
+ self.embed_dim = config.hidden_size
+ self.self_attn = CLIPSegAttention(config)
+ self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+ self.mlp = CLIPSegMLP(config)
+ self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.FloatTensor, torch.Tensor | None]:
+ residual = hidden_states
+
+ hidden_states = self.layer_norm1(hidden_states)
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.layer_norm2(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+class CLIPSegDecoderLayer(nn.Module):
+ """
+ CLIPSeg decoder layer, which is identical to `CLIPSegEncoderLayer`, except that normalization is applied after
+ self-attention/MLP, rather than before.
+ """
+
+ # Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer.__init__ with AltCLIP->CLIPSeg
+ def __init__(self, config: CLIPSegConfig):
+ super().__init__()
+ self.embed_dim = config.hidden_size
+ self.self_attn = CLIPSegAttention(config)
+ self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+ self.mlp = CLIPSegMLP(config)
+ self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ **kwargs,
+ ) -> tuple[torch.FloatTensor]:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ `(config.encoder_attention_heads,)`.
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+ returned tensors for more detail.
+ """
+ residual = hidden_states
+
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+
+ hidden_states = residual + hidden_states
+ hidden_states = self.layer_norm1(hidden_states)
+
+ residual = hidden_states
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+ hidden_states = self.layer_norm2(hidden_states)
+
+ return hidden_states
+
+
+@auto_docstring
+class CLIPSegPreTrainedModel(PreTrainedModel):
+ config: CLIPSegConfig
+ base_model_prefix = "clip"
+ input_modalities = ("image", "text")
+ supports_gradient_checkpointing = True
+ _can_record_outputs = {
+ "hidden_states": [CLIPSegEncoderLayer, CLIPSegDecoderLayer],
+ "attentions": CLIPSegAttention,
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ factor = self.config.initializer_factor
+ if isinstance(module, CLIPSegTextEmbeddings):
+ init.normal_(module.token_embedding.weight, mean=0.0, std=factor * 0.02)
+ init.normal_(module.position_embedding.weight, mean=0.0, std=factor * 0.02)
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+ elif isinstance(module, CLIPSegVisionEmbeddings):
+ factor = self.config.initializer_factor
+ init.normal_(module.class_embedding, mean=0.0, std=module.embed_dim**-0.5 * factor)
+ init.normal_(module.patch_embedding.weight, std=module.config.initializer_range * factor)
+ init.normal_(module.position_embedding.weight, std=module.config.initializer_range * factor)
+ init.copy_(module.position_ids, torch.arange(module.num_positions).expand((1, -1)))
+ elif isinstance(module, CLIPSegAttention):
+ factor = self.config.initializer_factor
+ in_proj_std = (module.embed_dim**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
+ out_proj_std = (module.embed_dim**-0.5) * factor
+ init.normal_(module.q_proj.weight, std=in_proj_std)
+ init.normal_(module.k_proj.weight, std=in_proj_std)
+ init.normal_(module.v_proj.weight, std=in_proj_std)
+ init.normal_(module.out_proj.weight, std=out_proj_std)
+ elif isinstance(module, CLIPSegMLP):
+ factor = self.config.initializer_factor
+ in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
+ fc_std = (2 * module.config.hidden_size) ** -0.5 * factor
+ init.normal_(module.fc1.weight, std=fc_std)
+ init.normal_(module.fc2.weight, std=in_proj_std)
+ elif isinstance(module, CLIPSegModel):
+ init.normal_(
+ module.text_projection.weight,
+ std=module.text_embed_dim**-0.5 * self.config.initializer_factor,
+ )
+ init.normal_(
+ module.visual_projection.weight,
+ std=module.vision_embed_dim**-0.5 * self.config.initializer_factor,
+ )
+
+ if isinstance(module, nn.LayerNorm):
+ init.zeros_(module.bias)
+ init.ones_(module.weight)
+ if isinstance(module, nn.Linear) and module.bias is not None:
+ init.zeros_(module.bias)
+
+
+# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->CLIPSeg
+class CLIPSegEncoder(nn.Module):
+ """
+ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+ [`CLIPSegEncoderLayer`].
+
+ Args:
+ config: CLIPSegConfig
+ """
+
+ def __init__(self, config: CLIPSegConfig):
+ super().__init__()
+ self.config = config
+ self.layers = nn.ModuleList([CLIPSegEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ inputs_embeds,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutput:
+ r"""
+ Args:
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ 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.
+ 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)
+
+ """
+ hidden_states = inputs_embeds
+ for encoder_layer in self.layers:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask,
+ **kwargs,
+ )
+
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+class CLIPSegDecoder(CLIPSegPreTrainedModel):
+ def __init__(self, config: CLIPSegConfig):
+ super().__init__(config)
+
+ self.conditional_layer = config.conditional_layer
+
+ self.film_mul = nn.Linear(config.projection_dim, config.reduce_dim)
+ self.film_add = nn.Linear(config.projection_dim, config.reduce_dim)
+
+ if config.use_complex_transposed_convolution:
+ transposed_kernels = (config.vision_config.patch_size // 4, config.vision_config.patch_size // 4)
+
+ self.transposed_convolution = nn.Sequential(
+ nn.Conv2d(config.reduce_dim, config.reduce_dim, kernel_size=3, padding=1),
+ nn.ReLU(),
+ nn.ConvTranspose2d(
+ config.reduce_dim,
+ config.reduce_dim // 2,
+ kernel_size=transposed_kernels[0],
+ stride=transposed_kernels[0],
+ ),
+ nn.ReLU(),
+ nn.ConvTranspose2d(
+ config.reduce_dim // 2, 1, kernel_size=transposed_kernels[1], stride=transposed_kernels[1]
+ ),
+ )
+ else:
+ self.transposed_convolution = nn.ConvTranspose2d(
+ config.reduce_dim, 1, config.vision_config.patch_size, stride=config.vision_config.patch_size
+ )
+
+ depth = len(config.extract_layers)
+ self.reduces = nn.ModuleList(
+ [nn.Linear(config.vision_config.hidden_size, config.reduce_dim) for _ in range(depth)]
+ )
+
+ decoder_config = copy.deepcopy(config.vision_config)
+ decoder_config.hidden_size = config.reduce_dim
+ decoder_config.num_attention_heads = config.decoder_num_attention_heads
+ decoder_config.intermediate_size = config.decoder_intermediate_size
+ decoder_config.hidden_act = "relu"
+ self.layers = nn.ModuleList([CLIPSegDecoderLayer(decoder_config) for _ in range(len(config.extract_layers))])
+
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(
+ self,
+ hidden_states: tuple[torch.Tensor],
+ conditional_embeddings: torch.Tensor,
+ output_attentions: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ activations = hidden_states[::-1]
+
+ output = None
+ for i, (activation, layer, reduce) in enumerate(zip(activations, self.layers, self.reduces)):
+ if output is not None:
+ output = reduce(activation) + output
+ else:
+ output = reduce(activation)
+
+ if i == self.conditional_layer:
+ output = self.film_mul(conditional_embeddings) * output.permute(1, 0, 2) + self.film_add(
+ conditional_embeddings
+ )
+ output = output.permute(1, 0, 2)
+
+ output = layer(output, attention_mask=None, causal_attention_mask=None, **kwargs)
+
+ output = output[:, 1:, :].permute(0, 2, 1) # remove cls token and reshape to [batch_size, reduce_dim, seq_len]
+
+ size = int(math.sqrt(output.shape[2]))
+
+ batch_size = conditional_embeddings.shape[0]
+ output = output.view(batch_size, output.shape[1], size, size)
+
+ logits = self.transposed_convolution(output).squeeze(1)
+
+ return CLIPSegDecoderOutput(logits=logits)
+
+
+class CLIPSegTextTransformer(CLIPSegPreTrainedModel):
+ def __init__(self, config: CLIPSegTextConfig):
+ super().__init__(config)
+
+ embed_dim = config.hidden_size
+ self.embeddings = CLIPSegTextEmbeddings(config)
+ self.encoder = CLIPSegEncoder(config)
+ self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+ # For `pooled_output` computation
+ self.eos_token_id = config.eos_token_id
+
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | BaseModelOutputWithPooling:
+ 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.return_dict
+
+ if input_ids is None:
+ raise ValueError("You have to specify input_ids")
+
+ input_shape = input_ids.size()
+ input_ids = input_ids.view(-1, input_shape[-1])
+
+ hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)
+
+ attention_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=hidden_states,
+ attention_mask=attention_mask,
+ past_key_values=None,
+ )
+
+ kwargs.pop("is_causal", None)
+ encoder_outputs = self.encoder(
+ inputs_embeds=hidden_states,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ is_causal=True,
+ **kwargs,
+ )
+
+ last_hidden_state = encoder_outputs[0]
+ last_hidden_state = self.final_layer_norm(last_hidden_state)
+
+ if self.eos_token_id == 2:
+ # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here.
+ # A CLIPSeg model with such `eos_token_id` in the config can't work correctly with extra new tokens added
+ # ------------------------------------------------------------
+ # text_embeds.shape = [batch_size, sequence_length, transformer.width]
+ # take features from the eot embedding (eot_token is the highest number in each sequence)
+ # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14
+ pooled_output = last_hidden_state[
+ torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
+ input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1),
+ ]
+ else:
+ # The config gets updated `eos_token_id` from PR #24773 (so the use of extra new tokens is possible)
+ pooled_output = last_hidden_state[
+ torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
+ # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`)
+ # Note: we assume each sequence (along batch dim.) contains an `eos_token_id` (e.g. prepared by the tokenizer)
+ (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.eos_token_id)
+ .int()
+ .argmax(dim=-1),
+ ]
+
+ if not return_dict:
+ return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ )
+
+
+class CLIPSegTextModel(CLIPSegPreTrainedModel):
+ config: CLIPSegTextConfig
+ input_modalities = ("text",)
+
+ _no_split_modules = ["CLIPSegTextEmbeddings", "CLIPSegEncoderLayer"]
+
+ def __init__(self, config: CLIPSegTextConfig):
+ super().__init__(config)
+ self.text_model = CLIPSegTextTransformer(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.text_model.embeddings.token_embedding
+
+ def set_input_embeddings(self, value):
+ self.text_model.embeddings.token_embedding = value
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, CLIPSegTextModel
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("CIDAS/clipseg-rd64-refined")
+ >>> model = CLIPSegTextModel.from_pretrained("CIDAS/clipseg-rd64-refined")
+
+ >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")
+
+ >>> outputs = model(**inputs)
+ >>> last_hidden_state = outputs.last_hidden_state
+ >>> pooled_output = outputs.pooler_output # pooled (EOS token) states
+ ```"""
+ return self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ return_dict=True,
+ **kwargs,
+ )
+
+
+class CLIPSegVisionTransformer(nn.Module):
+ # Copied from transformers.models.altclip.modeling_altclip.AltCLIPVisionTransformer.__init__ with AltCLIP->CLIPSeg
+ def __init__(self, config: CLIPSegVisionConfig):
+ super().__init__()
+ self.config = config
+ embed_dim = config.hidden_size
+
+ self.embeddings = CLIPSegVisionEmbeddings(config)
+ self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+ self.encoder = CLIPSegEncoder(config)
+ self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ interpolate_pos_encoding: bool | None = True,
+ ) -> tuple | BaseModelOutputWithPooling:
+ 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.return_dict
+
+ hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
+ hidden_states = self.pre_layrnorm(hidden_states)
+
+ encoder_outputs = self.encoder(
+ inputs_embeds=hidden_states,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ last_hidden_state = encoder_outputs[0]
+ pooled_output = last_hidden_state[:, 0, :]
+ pooled_output = self.post_layernorm(pooled_output)
+
+ if not return_dict:
+ return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ )
+
+
+class CLIPSegVisionModel(CLIPSegPreTrainedModel):
+ config: CLIPSegVisionConfig
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+
+ def __init__(self, config: CLIPSegVisionConfig):
+ super().__init__(config)
+ self.vision_model = CLIPSegVisionTransformer(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.vision_model.embeddings.patch_embedding
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ interpolate_pos_encoding: bool | None = True,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ Examples:
+
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoProcessor, CLIPSegVisionModel
+
+ >>> processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
+ >>> model = CLIPSegVisionModel.from_pretrained("CIDAS/clipseg-rd64-refined")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> inputs = processor(images=image, return_tensors="pt")
+
+ >>> outputs = model(**inputs)
+ >>> last_hidden_state = outputs.last_hidden_state
+ >>> pooled_output = outputs.pooler_output # pooled CLS states
+ ```"""
+ return self.vision_model(
+ pixel_values=pixel_values,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ **kwargs,
+ )
+
+
+@auto_docstring
+class CLIPSegModel(CLIPSegPreTrainedModel):
+ config: CLIPSegConfig
+
+ def __init__(self, config: CLIPSegConfig):
+ super().__init__(config)
+
+ if not isinstance(config.text_config, CLIPSegTextConfig):
+ raise TypeError(
+ "config.text_config is expected to be of type CLIPSegTextConfig but is of type"
+ f" {type(config.text_config)}."
+ )
+
+ if not isinstance(config.vision_config, CLIPSegVisionConfig):
+ raise TypeError(
+ "config.vision_config is expected to be of type CLIPSegVisionConfig but is of type"
+ f" {type(config.vision_config)}."
+ )
+
+ text_config = config.text_config
+ vision_config = config.vision_config
+ # The module using it is not a PreTrainedModel subclass so we need this
+ text_config._attn_implementation = config._attn_implementation
+ # The module using it is not a PreTrainedModel subclass so we need this
+ vision_config._attn_implementation = config._attn_implementation
+
+ self.projection_dim = config.projection_dim
+ self.text_embed_dim = text_config.hidden_size
+ self.vision_embed_dim = vision_config.hidden_size
+
+ self.text_model = CLIPSegTextTransformer(text_config)
+ self.vision_model = CLIPSegVisionTransformer(vision_config)
+
+ self.visual_projection = nn.Linear(self.vision_embed_dim, self.projection_dim, bias=False)
+ self.text_projection = nn.Linear(self.text_embed_dim, self.projection_dim, bias=False)
+ self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def get_text_features(
+ self,
+ input_ids: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ Examples:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoTokenizer, CLIPSegModel
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("CIDAS/clipseg-rd64-refined")
+ >>> model = CLIPSegModel.from_pretrained("CIDAS/clipseg-rd64-refined")
+
+ >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")
+ >>> with torch.inference_mode():
+ ... text_features = model.get_text_features(**inputs)
+ ```"""
+ text_outputs: BaseModelOutputWithPooling = self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ **kwargs,
+ )
+ pooled_output = text_outputs.pooler_output
+ text_outputs.pooler_output = self.text_projection(pooled_output)
+
+ return text_outputs
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ interpolate_pos_encoding: bool = True,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ Examples:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoProcessor, CLIPSegModel
+ >>> from transformers.image_utils import load_image
+
+ >>> processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
+ >>> model = CLIPSegModel.from_pretrained("CIDAS/clipseg-rd64-refined")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = load_image(url)
+
+ >>> inputs = processor(images=image, return_tensors="pt")
+
+ >>> with torch.inference_mode():
+ ... image_features = model.get_image_features(**inputs)
+ ```"""
+ vision_outputs: BaseModelOutputWithPooling = self.vision_model(
+ pixel_values=pixel_values,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ **kwargs,
+ )
+ pooled_output = vision_outputs.pooler_output
+ vision_outputs.pooler_output = self.visual_projection(pooled_output)
+
+ return vision_outputs
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ return_loss: bool | None = None,
+ interpolate_pos_encoding: bool = True,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CLIPSegOutput:
+ r"""
+ return_loss (`bool`, *optional*):
+ Whether or not to return the contrastive loss.
+
+ Examples:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoProcessor, CLIPSegModel
+ >>> from transformers.image_utils import load_image
+
+ >>> processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
+ >>> model = CLIPSegModel.from_pretrained("CIDAS/clipseg-rd64-refined")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = load_image(url)
+
+ >>> inputs = processor(
+ ... text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True
+ ... )
+
+ >>> with torch.inference_mode():
+ ... outputs = model(**inputs)
+ >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score
+ >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities
+ ```"""
+ vision_outputs = self.get_image_features(
+ pixel_values=pixel_values,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ **kwargs,
+ )
+
+ text_outputs = self.get_text_features(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+ image_embeds = vision_outputs.pooler_output
+
+ text_embeds = text_outputs.pooler_output
+
+ # normalized features
+ image_embeds = image_embeds / image_embeds.norm(p=2, dim=-1, keepdim=True)
+ text_embeds = text_embeds / text_embeds.norm(p=2, dim=-1, keepdim=True)
+
+ # cosine similarity as logits
+ logit_scale = self.logit_scale.exp()
+ logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
+ logits_per_image = logits_per_text.t()
+
+ loss = None
+ if return_loss:
+ loss = clipseg_loss(logits_per_text)
+
+ return CLIPSegOutput(
+ loss=loss,
+ logits_per_image=logits_per_image,
+ logits_per_text=logits_per_text,
+ text_embeds=text_embeds,
+ image_embeds=image_embeds,
+ text_model_output=text_outputs,
+ vision_model_output=vision_outputs,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ CLIPSeg model with a Transformer-based decoder on top for zero-shot and one-shot image segmentation.
+ """
+)
+class CLIPSegForImageSegmentation(CLIPSegPreTrainedModel):
+ config: CLIPSegConfig
+
+ def __init__(self, config: CLIPSegConfig):
+ super().__init__(config)
+
+ self.config = config
+
+ self.clip = CLIPSegModel(config)
+ self.extract_layers = config.extract_layers
+
+ self.decoder = CLIPSegDecoder(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_conditional_embeddings(
+ self,
+ batch_size: int | None = None,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ conditional_pixel_values: torch.Tensor | None = None,
+ ):
+ if input_ids is not None:
+ # compute conditional embeddings from texts
+ if len(input_ids) != batch_size:
+ raise ValueError("Make sure to pass as many prompt texts as there are query images")
+ with torch.no_grad():
+ conditional_embeddings = self.clip.get_text_features(
+ input_ids, attention_mask=attention_mask, position_ids=position_ids
+ ).pooler_output
+ elif conditional_pixel_values is not None:
+ # compute conditional embeddings from images
+ if len(conditional_pixel_values) != batch_size:
+ raise ValueError("Make sure to pass as many prompt images as there are query images")
+ with torch.no_grad():
+ conditional_embeddings = self.clip.get_image_features(conditional_pixel_values).pooler_output
+ else:
+ raise ValueError(
+ "Invalid conditional, should be either provided as `input_ids` or `conditional_pixel_values`"
+ )
+
+ return conditional_embeddings
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.FloatTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ conditional_pixel_values: torch.FloatTensor | None = None,
+ conditional_embeddings: torch.FloatTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ interpolate_pos_encoding: bool = True,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CLIPSegOutput:
+ r"""
+ conditional_pixel_values (`torch.FloatTensor`, *optional*):
+ The pixel values of the conditional images.
+ conditional_embeddings (`torch.FloatTensor` of shape `(batch_size, config.projection_dim)`, *optional*):
+ The conditional embeddings for the query images. If provided, the model will use this instead of computing
+ the embeddings from the conditional_pixel_values.
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+
+ Examples:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoProcessor, CLIPSegForImageSegmentation
+ >>> from transformers.image_utils import load_image
+
+ >>> processor = AutoProcessor.from_pretrained("CIDAS/clipseg-rd64-refined")
+ >>> model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = load_image(url)
+
+ >>> texts = ["a cat", "a remote", "a blanket"]
+ >>> inputs = processor(text=texts, images=[image] * len(texts), padding=True, return_tensors="pt")
+
+ >>> with torch.inference_mode():
+ ... outputs = model(**inputs)
+
+ >>> logits = outputs.logits
+ >>> print(logits.shape)
+ torch.Size([3, 352, 352])
+ ```"""
+ # step 1: forward the query images through the frozen CLIP vision encoder
+ with torch.no_grad():
+ kwargs["output_hidden_states"] = True # required to extract layers for the stages
+ vision_outputs = self.clip.get_image_features(
+ pixel_values=pixel_values,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ **kwargs,
+ )
+ pooled_output = vision_outputs.pooler_output
+
+ hidden_states = vision_outputs.hidden_states
+ # we add +1 here as the hidden states also include the initial embeddings
+ activations = [hidden_states[i + 1] for i in self.extract_layers]
+
+ # update vision_outputs
+ vision_outputs = BaseModelOutputWithPooling(
+ last_hidden_state=vision_outputs.last_hidden_state,
+ pooler_output=vision_outputs.pooler_output,
+ hidden_states=vision_outputs.hidden_states,
+ attentions=vision_outputs.attentions,
+ )
+
+ # step 2: compute conditional embeddings, either from text, images or an own provided embedding
+ if conditional_embeddings is None:
+ conditional_embeddings = self.get_conditional_embeddings(
+ batch_size=pixel_values.shape[0],
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ conditional_pixel_values=conditional_pixel_values,
+ )
+ else:
+ if conditional_embeddings.shape[0] != pixel_values.shape[0]:
+ raise ValueError(
+ "Make sure to pass as many conditional embeddings as there are query images in the batch"
+ )
+ if conditional_embeddings.shape[1] != self.config.projection_dim:
+ raise ValueError(
+ "Make sure that the feature dimension of the conditional embeddings matches"
+ " `config.projection_dim`."
+ )
+
+ # step 3: forward both the pooled output and the activations through the lightweight decoder to predict masks
+ decoder_outputs = self.decoder(
+ activations,
+ conditional_embeddings,
+ **kwargs,
+ )
+ logits = decoder_outputs.logits
+
+ loss = None
+ if labels is not None:
+ # move labels to the correct device to enable PP
+ labels = labels.to(logits.device)
+ loss_fn = nn.BCEWithLogitsLoss()
+ loss = loss_fn(logits, labels)
+
+ return CLIPSegImageSegmentationOutput(
+ loss=loss,
+ logits=logits,
+ conditional_embeddings=conditional_embeddings,
+ pooled_output=pooled_output,
+ vision_model_output=vision_outputs,
+ decoder_output=decoder_outputs,
+ )
+
+
+__all__ = [
+ "CLIPSegModel",
+ "CLIPSegPreTrainedModel",
+ "CLIPSegTextModel",
+ "CLIPSegVisionModel",
+ "CLIPSegForImageSegmentation",
+]
diff --git a/third_party/transformers/src/transformers/models/clipseg/processing_clipseg.py b/third_party/transformers/src/transformers/models/clipseg/processing_clipseg.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d2e0538401a4679540c386a42a6660907e4d956
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/clipseg/processing_clipseg.py
@@ -0,0 +1,88 @@
+# Copyright 2022 The HuggingFace Inc. team.
+#
+# 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.
+"""
+Image/Text processor class for CLIPSeg
+"""
+
+from ...processing_utils import ProcessorMixin
+from ...tokenization_utils_base import BatchEncoding
+from ...utils import auto_docstring
+
+
+@auto_docstring
+class CLIPSegProcessor(ProcessorMixin):
+ def __init__(self, image_processor=None, tokenizer=None, **kwargs):
+ super().__init__(image_processor, tokenizer)
+
+ @auto_docstring
+ def __call__(self, text=None, images=None, visual_prompt=None, return_tensors=None, **kwargs):
+ r"""
+ visual_prompt (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):
+ The visual prompt image or batch of images to be prepared. Each visual prompt image can be a PIL image,
+ NumPy array or PyTorch tensor. In case of a NumPy array/PyTorch tensor, each image should be of shape
+ (C, H, W), where C is a number of channels, H and W are image height and width.
+
+ Returns:
+ [`BatchEncoding`]: A [`BatchEncoding`] with the following fields:
+
+ - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+ - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+ `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+ `None`).
+ - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+ """
+ if text is None and visual_prompt is None and images is None:
+ raise ValueError("You have to specify either text, visual prompt or images.")
+
+ if text is not None and visual_prompt is not None:
+ raise ValueError("You have to specify exactly one type of prompt. Either text or visual prompt.")
+
+ output_kwargs = self._merge_kwargs(
+ self.valid_processor_kwargs, tokenizer_init_kwargs=self.tokenizer.init_kwargs, **kwargs
+ )
+
+ if text is not None:
+ encoding = self.tokenizer(text, return_tensors=return_tensors, **output_kwargs["text_kwargs"])
+
+ if visual_prompt is not None:
+ prompt_features = self.image_processor(
+ visual_prompt, return_tensors=return_tensors, **output_kwargs["images_kwargs"]
+ )
+
+ if images is not None:
+ image_features = self.image_processor(
+ images, return_tensors=return_tensors, **output_kwargs["images_kwargs"]
+ )
+
+ if visual_prompt is not None and images is not None:
+ encoding = {
+ "pixel_values": image_features.pixel_values,
+ "conditional_pixel_values": prompt_features.pixel_values,
+ }
+ return encoding
+ elif text is not None and images is not None:
+ encoding["pixel_values"] = image_features.pixel_values
+ return encoding
+ elif text is not None:
+ return encoding
+ elif visual_prompt is not None:
+ encoding = {
+ "conditional_pixel_values": prompt_features.pixel_values,
+ }
+ return encoding
+ else:
+ return BatchEncoding(data=dict(**image_features), tensor_type=return_tensors)
+
+
+__all__ = ["CLIPSegProcessor"]
diff --git a/third_party/transformers/src/transformers/models/codegen/__init__.py b/third_party/transformers/src/transformers/models/codegen/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..a06b0c4883f5e9d98469acdb0edad8f5b4b67fec
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/codegen/__init__.py
@@ -0,0 +1,29 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from ..gpt2.tokenization_gpt2 import GPT2Tokenizer as CodeGenTokenizerFast
+ from .configuration_codegen import *
+ from .modeling_codegen import *
+ from .tokenization_codegen import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/codegen/configuration_codegen.py b/third_party/transformers/src/transformers/models/codegen/configuration_codegen.py
new file mode 100644
index 0000000000000000000000000000000000000000..c4eb70bdf3d07e47a861e9a92880a7a5ad78e7ce
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/codegen/configuration_codegen.py
@@ -0,0 +1,77 @@
+# Copyright 2022 Salesforce authors, The EleutherAI, and HuggingFace Teams. All rights reserved.
+#
+# 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.
+"""CodeGen model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="Salesforce/codegen-2B-mono")
+@strict
+class CodeGenConfig(PreTrainedConfig):
+ r"""
+ n_ctx (`int`, *optional*, defaults to 2048):
+ This attribute is used in `CodeGenModel.__init__` without any real effect.
+ The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+ rotary_dim (`int`, *optional*, defaults to 64):
+ Number of dimensions in the embedding that Rotary Position Embedding is applied to.
+ n_inner (`int`, *optional*):
+ Dimensionality of the inner feed-forward layers. `None` will set it to 4 times n_embd
+
+ Example:
+
+ ```python
+ >>> from transformers import CodeGenConfig, CodeGenModel
+
+ >>> # Initializing a CodeGen 6B configuration
+ >>> configuration = CodeGenConfig()
+
+ >>> # Initializing a model (with random weights) from the configuration
+ >>> model = CodeGenModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "codegen"
+ attribute_map = {
+ "max_position_embeddings": "n_positions",
+ "hidden_size": "n_embd",
+ "num_attention_heads": "n_head",
+ "num_hidden_layers": "n_layer",
+ }
+
+ vocab_size: int = 50400
+ n_positions: int = 2048
+ n_ctx: int = 2048
+ n_embd: int = 4096
+ n_layer: int = 28
+ n_head: int = 16
+ rotary_dim: int = 64
+ n_inner: int | None = None
+ activation_function: str = "gelu_new"
+ resid_pdrop: float | int = 0.0
+ embd_pdrop: float | int = 0.0
+ attn_pdrop: float | int = 0.0
+ layer_norm_epsilon: float = 1e-5
+ initializer_range: float = 0.02
+ use_cache: bool = True
+ bos_token_id: int | None = 50256
+ eos_token_id: int | list[int] | None = 50256
+ tie_word_embeddings: bool = False
+
+
+__all__ = ["CodeGenConfig"]
diff --git a/third_party/transformers/src/transformers/models/codegen/modeling_codegen.py b/third_party/transformers/src/transformers/models/codegen/modeling_codegen.py
new file mode 100644
index 0000000000000000000000000000000000000000..d0fe74dd8714d81e9c7bb9943dde0792e4ed0571
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/codegen/modeling_codegen.py
@@ -0,0 +1,486 @@
+# Copyright 2022 Salesforce authors, The EleutherAI, and HuggingFace Teams. All rights reserved.
+#
+# 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 CodeGen model."""
+
+import math
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
+ auto_docstring,
+ logging,
+)
+from .configuration_codegen import CodeGenConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from transformers.models.gptj.modeling_gptj.create_sinusoidal_positions
+def create_sinusoidal_positions(num_pos: int, dim: int) -> torch.Tensor:
+ inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.int64) / dim))
+ sinusoid_inp = torch.einsum("i , j -> i j", torch.arange(num_pos, dtype=torch.int64).float(), inv_freq).float()
+ return torch.cat((torch.sin(sinusoid_inp), torch.cos(sinusoid_inp)), dim=1)
+
+
+# Copied from transformers.models.gptj.modeling_gptj.rotate_every_two
+def rotate_every_two(x: torch.Tensor) -> torch.Tensor:
+ x1 = x[:, :, :, ::2]
+ x2 = x[:, :, :, 1::2]
+ x = torch.stack((-x2, x1), dim=-1)
+ return x.flatten(-2) # in einsum notation: rearrange(x, '... d j -> ... (d j)')
+
+
+# Copied from transformers.models.gptj.modeling_gptj.apply_rotary_pos_emb
+def apply_rotary_pos_emb(tensor: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor) -> torch.Tensor:
+ sin = torch.repeat_interleave(sin[:, :, None, :], 2, 3)
+ cos = torch.repeat_interleave(cos[:, :, None, :], 2, 3)
+ return (tensor * cos) + (rotate_every_two(tensor) * sin)
+
+
+class CodeGenAttention(nn.Module):
+ def __init__(self, config, layer_idx=None):
+ super().__init__()
+
+ self.max_positions = config.max_position_embeddings
+ self.attn_dropout = nn.Dropout(config.attn_pdrop)
+ self.resid_dropout = nn.Dropout(config.resid_pdrop)
+ self.layer_idx = layer_idx
+ if layer_idx is None:
+ logger.warning_once(
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class."
+ )
+
+ self.embed_dim = config.hidden_size
+ self.num_attention_heads = config.num_attention_heads
+ self.head_dim = self.embed_dim // self.num_attention_heads
+ if self.head_dim * self.num_attention_heads != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_attention_heads (got `embed_dim`: {self.embed_dim} and"
+ f" `num_attention_heads`: {self.num_attention_heads})."
+ )
+ self.scale_attn = math.sqrt(self.head_dim)
+ self.qkv_proj = nn.Linear(self.embed_dim, self.embed_dim * 3, bias=False)
+
+ self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=False)
+ self.rotary_dim = config.rotary_dim
+ self.pos_embd_dim = self.rotary_dim or self.embed_dim
+ self.register_buffer(
+ "embed_positions", create_sinusoidal_positions(self.max_positions, self.pos_embd_dim), persistent=False
+ )
+
+ def _split_heads(self, x, n_head, dim_head, mp_num):
+ reshaped = x.reshape(x.shape[:-1] + (n_head // mp_num, dim_head))
+ reshaped = reshaped.reshape(x.shape[:-2] + (-1,) + reshaped.shape[-1:])
+ return reshaped
+
+ def _merge_heads(self, tensor, num_attention_heads, attn_head_size):
+ """
+ Merges attn_head_size dim and num_attn_heads dim into n_ctx
+ """
+ if len(tensor.shape) == 5:
+ tensor = tensor.permute(0, 1, 3, 2, 4).contiguous()
+ elif len(tensor.shape) == 4:
+ tensor = tensor.permute(0, 2, 1, 3).contiguous()
+ else:
+ raise ValueError(f"Input tensor rank should be one of [4, 5], but is: {len(tensor.shape)}")
+ new_shape = tensor.size()[:-2] + (num_attention_heads * attn_head_size,)
+ return tensor.view(new_shape)
+
+ def _attn(
+ self,
+ query,
+ key,
+ value,
+ attention_mask=None,
+ ):
+ # Keep the attention weights computation in fp32 to avoid overflow issues
+ query = query.to(torch.float32)
+ key = key.to(torch.float32)
+
+ attn_weights = torch.matmul(query, key.transpose(-1, -2))
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = attn_weights / self.scale_attn
+ attn_weights = nn.Softmax(dim=-1)(attn_weights)
+ attn_weights = attn_weights.to(value.dtype)
+ attn_weights = self.attn_dropout(attn_weights)
+
+ attn_output = torch.matmul(attn_weights, value)
+
+ return attn_output, attn_weights
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor | None,
+ layer_past: Cache | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ use_cache: bool | None = False,
+ output_attentions: bool | None = False,
+ ) -> (
+ tuple[torch.Tensor, tuple[torch.Tensor]]
+ | tuple[torch.Tensor, tuple[torch.Tensor], tuple[torch.Tensor, ...]]
+ | None
+ ):
+ qkv = self.qkv_proj(hidden_states)
+ # TODO(enijkamp): factor out number of logical TPU-v4 cores or make forward pass agnostic
+ mp_num = 4
+ qkv_split = qkv.reshape(qkv.shape[:-1] + (mp_num, -1))
+
+ local_dim = self.head_dim * self.num_attention_heads // mp_num
+ query, value, key = torch.split(qkv_split, local_dim, dim=-1)
+ query = self._split_heads(query, self.num_attention_heads, self.head_dim, mp_num=mp_num)
+ key = self._split_heads(key, self.num_attention_heads, self.head_dim, mp_num=mp_num)
+
+ value = self._split_heads(value, self.num_attention_heads, self.head_dim, mp_num=mp_num)
+ value = value.permute(0, 2, 1, 3)
+
+ embed_positions = self.embed_positions
+ if embed_positions.device != position_ids.device:
+ embed_positions = embed_positions.to(position_ids.device)
+ self.embed_positions = embed_positions
+
+ sincos = embed_positions[position_ids]
+ sin, cos = torch.split(sincos, sincos.shape[-1] // 2, dim=-1)
+
+ if self.rotary_dim is not None:
+ k_rot = key[:, :, :, : self.rotary_dim]
+ k_pass = key[:, :, :, self.rotary_dim :]
+
+ q_rot = query[:, :, :, : self.rotary_dim]
+ q_pass = query[:, :, :, self.rotary_dim :]
+
+ k_rot = apply_rotary_pos_emb(k_rot, sin, cos)
+ q_rot = apply_rotary_pos_emb(q_rot, sin, cos)
+
+ key = torch.cat([k_rot, k_pass], dim=-1)
+ query = torch.cat([q_rot, q_pass], dim=-1)
+ else:
+ key = apply_rotary_pos_emb(key, sin, cos)
+ query = apply_rotary_pos_emb(query, sin, cos)
+
+ key = key.permute(0, 2, 1, 3)
+ query = query.permute(0, 2, 1, 3)
+
+ # Note that this cast is quite ugly, but is not implemented before ROPE as k_rot in the original codebase is always in fp32.
+ # Reference: https://github.com/salesforce/CodeGen/blob/f210c3bb1216c975ad858cd4132c0fdeabf4bfc2/codegen1/jaxformer/hf/codegen/modeling_codegen.py#L38
+ if layer_past is not None:
+ key, value = layer_past.update(key.to(hidden_states.dtype), value, self.layer_idx)
+
+ # compute self-attention: V x Softmax(QK^T)
+ attn_output, attn_weights = self._attn(query, key, value, attention_mask)
+
+ attn_output = self._merge_heads(attn_output, self.num_attention_heads, self.head_dim)
+ attn_output = self.out_proj(attn_output)
+ attn_output = self.resid_dropout(attn_output)
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.gptj.modeling_gptj.GPTJMLP with GPTJ->CodeGen
+class CodeGenMLP(nn.Module):
+ def __init__(self, intermediate_size, config): # in MLP: intermediate_size= 4 * embed_dim
+ super().__init__()
+ embed_dim = config.n_embd
+
+ self.fc_in = nn.Linear(embed_dim, intermediate_size)
+ self.fc_out = nn.Linear(intermediate_size, embed_dim)
+
+ self.act = ACT2FN[config.activation_function]
+ self.dropout = nn.Dropout(config.resid_pdrop)
+
+ def forward(self, hidden_states: torch.FloatTensor | None) -> torch.FloatTensor:
+ hidden_states = self.fc_in(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.fc_out(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.gptj.modeling_gptj.GPTJBlock with GPTJ->CodeGen
+class CodeGenBlock(GradientCheckpointingLayer):
+ # Ignore copy
+ def __init__(self, config, layer_idx=None):
+ super().__init__()
+ inner_dim = config.n_inner if config.n_inner is not None else 4 * config.n_embd
+ self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+ self.attn = CodeGenAttention(config, layer_idx)
+ self.mlp = CodeGenMLP(inner_dim, config)
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor | None,
+ layer_past: Cache | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ use_cache: bool | None = False,
+ output_attentions: bool | None = False,
+ **kwargs,
+ ) -> tuple[torch.Tensor] | tuple[torch.Tensor, tuple[torch.FloatTensor, ...]] | None:
+ residual = hidden_states
+ hidden_states = self.ln_1(hidden_states)
+ attn_outputs, attn_weights = self.attn(
+ hidden_states=hidden_states,
+ layer_past=layer_past,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ )
+ feed_forward_hidden_states = self.mlp(hidden_states)
+ hidden_states = attn_outputs + feed_forward_hidden_states + residual
+
+ return hidden_states, attn_weights
+
+
+@auto_docstring
+class CodeGenPreTrainedModel(PreTrainedModel):
+ config: CodeGenConfig
+ base_model_prefix = "transformer"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["CodeGenBlock"]
+ _skip_keys_device_placement = "past_key_values"
+ _can_compile_fullgraph = True
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, CodeGenAttention):
+ init.copy_(module.embed_positions, create_sinusoidal_positions(module.max_positions, module.pos_embd_dim))
+
+
+@auto_docstring
+class CodeGenModel(CodeGenPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.embed_dim = config.n_embd
+ self.vocab_size = config.vocab_size
+ self.wte = nn.Embedding(config.vocab_size, self.embed_dim)
+ self.drop = nn.Dropout(config.embd_pdrop)
+ self.h = nn.ModuleList([CodeGenBlock(config, layer_idx=i) for i in range(config.n_layer)])
+ self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
+ self.rotary_dim = min(config.rotary_dim, config.n_ctx // config.num_attention_heads)
+
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.wte
+
+ def set_input_embeddings(self, new_embeddings):
+ self.wte = new_embeddings
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs, # NOOP kwargs, for now
+ ) -> tuple | BaseModelOutputWithPast:
+ r"""
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_dim)`, *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.
+ """
+ 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.return_dict
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or 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
+
+ if inputs_embeds is None:
+ inputs_embeds = self.wte(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ seq_length = inputs_embeds.shape[1]
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ hidden_states = inputs_embeds
+
+ if token_type_ids is not None:
+ token_type_ids = token_type_ids.view(-1, seq_length)
+ token_type_embeds = self.wte(token_type_ids)
+ hidden_states = hidden_states + token_type_embeds
+
+ hidden_states = self.drop(hidden_states)
+ output_shape = (-1, seq_length, hidden_states.size(-1))
+
+ all_self_attentions = () if output_attentions else None
+ all_hidden_states = () if output_hidden_states else None
+ for i, block in enumerate(self.h):
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ outputs = block(
+ hidden_states,
+ layer_past=past_key_values,
+ attention_mask=causal_mask,
+ position_ids=position_ids,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ )
+
+ hidden_states = outputs[0]
+ if output_attentions:
+ all_self_attentions = all_self_attentions + (outputs[1],)
+
+ hidden_states = self.ln_f(hidden_states)
+
+ hidden_states = hidden_states.view(output_shape)
+ # Add last hidden state
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ if not return_dict:
+ return tuple(
+ v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attentions] if v is not None
+ )
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The CodeGen Model transformer with a language modeling head on top.
+ """
+)
+class CodeGenForCausalLM(CodeGenPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "transformer.wte.weight"}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.transformer = CodeGenModel(config)
+ self.lm_head = nn.Linear(config.n_embd, config.vocab_size)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs,
+ ) -> tuple | CausalLMOutputWithPast:
+ r"""
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_dim)`, *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.
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+ `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+ are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ transformer_outputs = self.transformer(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = transformer_outputs[0]
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ if not return_dict:
+ output = (logits,) + transformer_outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=transformer_outputs.past_key_values,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+
+__all__ = ["CodeGenForCausalLM", "CodeGenModel", "CodeGenPreTrainedModel"]
diff --git a/third_party/transformers/src/transformers/models/codegen/tokenization_codegen.py b/third_party/transformers/src/transformers/models/codegen/tokenization_codegen.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d54a80592b5395fa63d4122cb15c3bbe33dcd1d
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/codegen/tokenization_codegen.py
@@ -0,0 +1,215 @@
+# Copyright 2022 The Salesforce authors, The Open AI Team Authors and The HuggingFace Inc. team.
+#
+# 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.
+"""Tokenization classes for CodeGen."""
+
+import re
+from typing import TYPE_CHECKING, Union
+
+import numpy as np
+from tokenizers import Tokenizer, decoders, pre_tokenizers, processors
+from tokenizers.models import BPE
+
+from ...tokenization_utils_tokenizers import TokenizersBackend
+from ...utils import is_torch_available, logging
+
+
+if TYPE_CHECKING:
+ if is_torch_available():
+ import torch
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt"}
+
+
+class CodeGenTokenizer(TokenizersBackend):
+ """
+ Construct a CodeGen tokenizer (backed by HuggingFace's *tokenizers* library). Based on byte-level
+ Byte-Pair-Encoding.
+
+ This tokenizer has been trained to treat spaces like parts of the tokens (a bit like sentencepiece) so a word will
+ be encoded differently whether it is at the beginning of the sentence (without space) or not:
+
+ ```python
+ >>> from transformers import CodeGenTokenizer
+
+ >>> tokenizer = CodeGenTokenizer.from_pretrained("Salesforce/codegen-350M-mono")
+ >>> tokenizer("Hello world")["input_ids"]
+ [15496, 995]
+
+ >>> tokenizer(" Hello world")["input_ids"]
+ [18435, 995]
+ ```
+
+ You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer, but since
+ the model was not pretrained this way, it might yield a decrease in performance.
+
+
+
+ When used with `is_split_into_words=True`, this tokenizer needs to be instantiated with `add_prefix_space=True`.
+
+
+
+ This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should
+ refer to this superclass for more information regarding those methods.
+
+ Args:
+ vocab (`str` or `dict[str, int]`, *optional*):
+ Custom vocabulary dictionary. If not provided, vocabulary is loaded from `vocab_file`.
+ merges (`str` or `list[str]`, *optional*):
+ Custom merges list. If not provided, merges are loaded from `merges_file`.
+ unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ bos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+ The beginning of sequence token.
+ eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+ The end of sequence token.
+ pad_token (`str`, *optional*):
+ The token used for padding, for example when batching sequences of different lengths.
+ add_prefix_space (`bool`, *optional*, defaults to `False`):
+ Whether or not to add an initial space to the input. This allows to treat the leading word just as any
+ other word. (CodeGen tokenizer detect beginning of words by the preceding space).
+ return_token_type_ids (`bool`, *optional*, defaults to `False`):
+ Whether to return token type IDs.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ model_input_names = ["input_ids", "attention_mask"]
+ model = BPE
+
+ def __init__(
+ self,
+ vocab: str | dict[str, int] | None = None,
+ merges: str | list[str] | None = None,
+ unk_token: str = "<|endoftext|>",
+ bos_token: str = "<|endoftext|>",
+ eos_token: str = "<|endoftext|>",
+ pad_token=None,
+ add_prefix_space: bool = False,
+ return_token_type_ids: bool = False,
+ **kwargs,
+ ):
+ self.return_token_type_ids = return_token_type_ids
+ if self.return_token_type_ids:
+ self.model_input_names.append("token_type_ids")
+
+ self.add_prefix_space = add_prefix_space
+
+ self._vocab = vocab if vocab is not None else {}
+ self._merges = merges or []
+
+ self._tokenizer = Tokenizer(
+ BPE(
+ vocab=self._vocab,
+ merges=self._merges,
+ dropout=None,
+ continuing_subword_prefix="",
+ end_of_word_suffix="",
+ fuse_unk=False,
+ )
+ )
+
+ self._tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=add_prefix_space)
+ self._tokenizer.decoder = decoders.ByteLevel()
+ self._tokenizer.post_processor = processors.ByteLevel(
+ add_prefix_space=True, use_regex=True, trim_offsets=False
+ )
+
+ super().__init__(
+ unk_token=unk_token,
+ bos_token=bos_token,
+ eos_token=eos_token,
+ pad_token=pad_token,
+ add_prefix_space=add_prefix_space,
+ return_token_type_ids=return_token_type_ids,
+ **kwargs,
+ )
+
+ def decode(
+ self,
+ token_ids: Union[int, list[int], np.ndarray, "torch.Tensor"],
+ skip_special_tokens: bool = False,
+ clean_up_tokenization_spaces: bool | None = None,
+ truncate_before_pattern: list[str] | None = None,
+ **kwargs,
+ ) -> str:
+ """
+ Converts a sequence of ids in a string, using the tokenizer and vocabulary with options to remove special
+ tokens and clean up tokenization spaces.
+
+ Similar to doing `self.convert_tokens_to_string(self.convert_ids_to_tokens(token_ids))`.
+
+ Args:
+ token_ids (`Union[int, List[int], np.ndarray, torch.Tensor]`):
+ List of tokenized input ids. Can be obtained using the `__call__` method.
+ skip_special_tokens (`bool`, *optional*, defaults to `False`):
+ Whether or not to remove special tokens in the decoding.
+ clean_up_tokenization_spaces (`bool`, *optional*):
+ Whether or not to clean up the tokenization spaces. If `None`, will default to
+ `self.clean_up_tokenization_spaces` (available in the `tokenizer_config`).
+ truncate_before_pattern (`List[str]`, *optional*, defaults to `None`):
+ A list of regular expression strings that will be used to truncate the returned string. This can be
+ used to remove extra pieces of code (e.g. truncate if observing a comment symbol "#" at the beginning
+ of a new line). An example pattern could be `["^#", re.escape("<|endoftext|>"), "^'''", "\n\n\n"]`.
+ kwargs (additional keyword arguments, *optional*):
+ Will be passed to the underlying model specific decode method.
+
+ Returns:
+ `str`: The decoded sentence.
+ """
+
+ decoded_text = super().decode(
+ token_ids=token_ids,
+ skip_special_tokens=skip_special_tokens,
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+ **kwargs,
+ )
+
+ if truncate_before_pattern is not None and len(truncate_before_pattern) > 0:
+ decoded_text = self.truncate(decoded_text, truncate_before_pattern)
+
+ return decoded_text
+
+ def truncate(self, completion, truncate_before_pattern):
+ def find_re(string, pattern, start_pos):
+ m = pattern.search(string, start_pos)
+ return m.start() if m else -1
+
+ terminals = [re.compile(pattern, re.MULTILINE) for pattern in truncate_before_pattern]
+
+ prints = list(re.finditer("^print", completion, re.MULTILINE))
+
+ if len(prints) > 1:
+ completion = completion[: prints[1].start()]
+
+ defs = list(re.finditer("^def", completion, re.MULTILINE))
+
+ if len(defs) > 1:
+ completion = completion[: defs[1].start()]
+
+ start_pos = 0
+
+ terminals_pos = [
+ pos for pos in [find_re(completion, terminal, start_pos) for terminal in terminals] if pos != -1
+ ]
+
+ if len(terminals_pos) > 0:
+ return completion[: min(terminals_pos)]
+ else:
+ return completion
+
+
+__all__ = ["CodeGenTokenizer"]
diff --git a/third_party/transformers/src/transformers/models/colqwen2/__init__.py b/third_party/transformers/src/transformers/models/colqwen2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d313293c0a0a7b5c56ac5bec1e54e8eb1e5a6e4a
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/colqwen2/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_colqwen2 import *
+ from .modeling_colqwen2 import *
+ from .processing_colqwen2 import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/colqwen2/configuration_colqwen2.py b/third_party/transformers/src/transformers/models/colqwen2/configuration_colqwen2.py
new file mode 100644
index 0000000000000000000000000000000000000000..ac9abcb5cfd91f0688bf6402b988207608f90d1c
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/colqwen2/configuration_colqwen2.py
@@ -0,0 +1,65 @@
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# 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.
+
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+from ..auto import CONFIG_MAPPING
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="vidore/colqwen2-v1.0-hf")
+@strict
+class ColQwen2Config(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ from transformers.models.colqwen2 import ColQwen2Config, ColQwen2ForRetrieval
+
+ config = ColQwen2Config()
+ model = ColQwen2ForRetrieval(config)
+ ```
+ """
+
+ model_type = "colqwen2"
+ sub_configs = {"vlm_config": PreTrainedConfig}
+
+ vlm_config: dict | PreTrainedConfig | None = None
+ embedding_dim: int = 128
+ initializer_range: float = 0.02
+
+ def __post_init__(self, **kwargs):
+ if self.vlm_config is None:
+ self.vlm_config = CONFIG_MAPPING["qwen2_vl"]()
+ logger.info(
+ "`vlm_config` is `None`. Initializing `vlm_config` with the `Qwen2VLConfig` with default values."
+ )
+ elif isinstance(self.vlm_config, dict):
+ self.vlm_config = CONFIG_MAPPING[self.vlm_config["model_type"]](**self.vlm_config)
+
+ if not hasattr(self.vlm_config, "vocab_size"):
+ self.vlm_config.vocab_size = self.vlm_config.get_text_config().vocab_size
+
+ super().__post_init__(**kwargs)
+
+ def get_text_config(self, *args, **kwargs) -> PreTrainedConfig:
+ return self.vlm_config.get_text_config(*args, **kwargs)
+
+
+__all__ = ["ColQwen2Config"]
diff --git a/third_party/transformers/src/transformers/models/colqwen2/convert_colqwen2_weights_to_hf.py b/third_party/transformers/src/transformers/models/colqwen2/convert_colqwen2_weights_to_hf.py
new file mode 100644
index 0000000000000000000000000000000000000000..496467c7b97d4bcc2781b2f20161c9f29ee01270
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/colqwen2/convert_colqwen2_weights_to_hf.py
@@ -0,0 +1,227 @@
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# 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.
+"""
+Convert ColQwen2 weights from the original repository to the HF model format.
+
+Don't forget to manually upload the processor-related files to the HF model repository
+after running this script.
+
+Original repository: https://github.com/illuin-tech/colqwen2.
+
+NOTE: This script was originally run using `torch==2.5.1` and with:
+
+```bash
+python src/transformers/models/colqwen2/convert_colqwen2_weights_to_hf.py \
+ --model_id vidore/colqwen2-v1.0-merged \
+ --revision eeccbae1d44bdcb0c83b1788127a2b2cad7d718e \
+ --original_vlm_name_or_path Qwen/Qwen2-VL-2B-Instruct \
+ --output_dir vidore/colqwen2-v1.0-hf-internal \
+ --push_to_hub
+```
+"""
+
+import argparse
+import glob
+from pathlib import Path
+from typing import Any
+
+import torch
+from huggingface_hub import snapshot_download
+from peft import PeftModel
+from safetensors import safe_open
+
+from transformers import AutoConfig, AutoModel
+from transformers.models.colqwen2 import ColQwen2ForRetrieval
+from transformers.models.colqwen2.configuration_colqwen2 import ColQwen2Config
+from transformers.utils import logging
+
+
+logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+
+ORIGINAL_DTYPE = torch.bfloat16
+
+
+def load_original_state_dict(model_id: str, revision: str | None = None) -> dict[str, torch.Tensor]:
+ directory_path = snapshot_download(
+ repo_id=model_id,
+ revision=revision,
+ allow_patterns=["*.safetensors"],
+ )
+
+ original_state_dict = {}
+ for path in glob.glob(f"{directory_path}/*"):
+ if path.endswith(".safetensors"):
+ with safe_open(path, framework="pt", device="cpu") as f:
+ for key in f.keys():
+ original_state_dict[key] = f.get_tensor(key)
+
+ # Some weights are tied, so `lm.head`` is not saved. Let's clone to load state dict.
+ if "lm_head.weight" not in original_state_dict and "model.embed_tokens.weight" in original_state_dict:
+ original_state_dict["lm_head.weight"] = original_state_dict["model.embed_tokens.weight"].clone()
+
+ return original_state_dict
+
+
+def rename_state_dict_keys(state_dict: dict[str, Any]) -> dict[str, Any]:
+ new_state_dict: dict[str, Any] = {}
+ for key, value in state_dict.items():
+ if key.startswith("custom_text_proj"):
+ new_key = key.replace("custom_text_proj", "embedding_proj_layer")
+ else:
+ # The original ColQwen2 inherits from Qwen2VL, so we simply need to add the `vlm.` prefix
+ # to all remaining keys.
+ if key.startswith("model."):
+ key = key.replace("model.", "model.language_model.")
+ if key.startswith("visual."):
+ key = key.replace("visual.", "model.visual.")
+ new_key = "vlm." + key
+ new_state_dict[new_key] = value
+ return new_state_dict
+
+
+@torch.no_grad()
+def convert_colqwen2_weights_to_hf(
+ model_id: str,
+ output_dir: str,
+ push_to_hub: bool,
+ revision: str | None = None,
+ original_vlm_name_or_path: str | None = None,
+):
+ # Load the original model data
+ original_config = AutoConfig.from_pretrained(
+ model_id,
+ revision=revision,
+ )
+ if original_vlm_name_or_path is not None:
+ original_config._name_or_path = original_vlm_name_or_path
+ if hasattr(original_config, "architectures"):
+ delattr(original_config, "architectures")
+
+ original_state_dict = load_original_state_dict(model_id, revision=revision)
+
+ # Format the state_dict keys
+ original_state_dict = rename_state_dict_keys(original_state_dict)
+
+ # Create the new config
+ config = ColQwen2Config(
+ vlm_config=original_config,
+ embedding_dim=128, # hardcoded in the original model
+ )
+ config.model_type = "colqwen2"
+ config.is_composition = False
+
+ # Load the untrained model
+ vlm_name_or_path = getattr(config.vlm_config, "_name_or_path", None)
+ if vlm_name_or_path and "2.5" in str(vlm_name_or_path):
+ print(
+ "Detected colqwen2.5 adapters in vlm_config; loading base model %s and merging PEFT weights."
+ % vlm_name_or_path
+ )
+ base_model = AutoModel.from_pretrained(
+ vlm_name_or_path,
+ device_map="cpu",
+ trust_remote_code=True,
+ )
+ peft_model = PeftModel.from_pretrained(base_model, model_id)
+ model = peft_model.merge_and_unload()
+ else:
+ model = ColQwen2ForRetrieval(config=config).to("cpu").eval()
+ print("Created model with new config and randomly initialized weights")
+
+ # NOTE: The new model was initialized with float32 weights. We need to convert it to the desired precision.
+ # There are two ways to set the model's dtype:
+ # - Using `model.from_pretrained(..., dtype=dtype_precision)` doesn't convert the hyperparameters to the desired precision.
+ # - Using `model.to(dtype_precision)` converts all values - including the hyperparameters - to the desired precision.
+ # The following snippet allows a fine-grained control over the model's dtype, making sure that all
+ # the new weights' dtypes match the original model.
+ for param in model.parameters():
+ param.data = param.data.to(ORIGINAL_DTYPE)
+ print(f"Converted the new model weights to `{ORIGINAL_DTYPE}`")
+
+ # Load the original weights
+ model.load_state_dict(original_state_dict)
+ print("Loaded original model weights")
+
+ # # Sanity check: ensure all keys are the same
+ state_dict_keys_old = set(original_state_dict.keys())
+ state_dict_keys_new = set(model.state_dict().keys())
+ disjoint_keys = state_dict_keys_old.symmetric_difference(state_dict_keys_new)
+ if disjoint_keys:
+ raise ValueError(f"Incompatible keys: {disjoint_keys}")
+
+ # Save the model
+ if push_to_hub:
+ model.push_to_hub(output_dir, private=True)
+ print(f"Model pushed to the hub at `{output_dir}`")
+ else:
+ Path(output_dir).mkdir(exist_ok=True, parents=True)
+ model.save_pretrained(output_dir)
+ print(f"Model saved to `{output_dir}`")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser(
+ description="""
+ This script converts the original ColQwen2 model to the HF model format.
+
+ Don't forget to manually upload the processor-related files to the HF model repository
+ after running this script.
+
+ Example usage:
+ ```bash
+ python src/transformers/models/colqwen2/convert_colqwen2_weights_to_hf.py \
+ --model_id vidore/colqwen2-v1.0-merged \
+ --revision eeccbae1d44bdcb0c83b1788127a2b2cad7d718e \
+ --original_vlm_name_or_path Qwen/Qwen2-VL-2B-Instruct \
+ --output_dir vidore/colqwen2-v1.0-hf-internal \
+ --push_to_hub
+ ```
+ """
+ )
+ parser.add_argument(
+ "--model_id",
+ help="Model ID of the original model to convert",
+ )
+ parser.add_argument(
+ "--output_dir",
+ help="Location to write HF model and tokenizer",
+ )
+ parser.add_argument(
+ "--push_to_hub",
+ help="Whether or not to push the model to the hub at `output_dir` instead of saving it locally",
+ action="store_true",
+ default=False,
+ )
+ parser.add_argument(
+ "--revision",
+ help="Revision of the model to download",
+ default=None,
+ )
+ parser.add_argument(
+ "--original_vlm_name_or_path",
+ help="Name or path of the original VLM backbone model",
+ default=None,
+ )
+
+ args = parser.parse_args()
+
+ convert_colqwen2_weights_to_hf(
+ model_id=args.model_id,
+ output_dir=args.output_dir,
+ push_to_hub=args.push_to_hub,
+ revision=args.revision,
+ original_vlm_name_or_path=args.original_vlm_name_or_path,
+ )
diff --git a/third_party/transformers/src/transformers/models/colqwen2/modeling_colqwen2.py b/third_party/transformers/src/transformers/models/colqwen2/modeling_colqwen2.py
new file mode 100644
index 0000000000000000000000000000000000000000..656ad6c758c5cb888e0a80589572bdd4fcc87389
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/colqwen2/modeling_colqwen2.py
@@ -0,0 +1,205 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/colqwen2/modular_colqwen2.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_colqwen2.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# 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.
+
+from dataclasses import dataclass
+
+from torch import nn
+
+from transformers import AutoModel
+
+from ... import initialization as init
+from ...cache_utils import Cache
+from ...modeling_utils import PreTrainedModel
+from ...utils import ModelOutput, auto_docstring, can_return_tuple, is_torch_available
+from .configuration_colqwen2 import ColQwen2Config
+
+
+if is_torch_available():
+ import torch
+
+
+@auto_docstring
+class ColQwen2PreTrainedModel(PreTrainedModel):
+ config: ColQwen2Config
+ base_model_prefix = "model"
+ input_modalities = ("image", "text")
+ _no_split_modules = []
+ _supports_sdpa = True
+ _supports_flash_attn = True
+ _supports_flex_attn = True
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ std = (
+ self.config.initializer_range
+ if hasattr(self.config, "initializer_range")
+ else self.config.vlm_config.text_config.initializer_range
+ )
+
+ if isinstance(module, (nn.Linear, nn.Conv2d)):
+ init.normal_(module.weight, mean=0.0, std=std)
+ if module.bias is not None:
+ init.zeros_(module.bias)
+ elif isinstance(module, nn.Embedding):
+ init.normal_(module.weight, mean=0.0, std=std)
+ # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag
+ if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False):
+ init.zeros_(module.weight[module.padding_idx])
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for ColQwen2 embeddings output.
+ """
+)
+class ColQwen2ForRetrievalOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token prediction).
+ embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ The embeddings of the model.
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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.
+ """
+
+ loss: torch.FloatTensor | None = None
+ embeddings: torch.Tensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Following the ColPali approach, ColQwen2 leverages VLMs to construct efficient multi-vector embeddings directly
+ from document images (“screenshots”) for document retrieval. The model is trained to maximize the similarity
+ between these document embeddings and the corresponding query embeddings, using the late interaction method
+ introduced in ColBERT.
+
+ Using ColQwen2 removes the need for potentially complex and brittle layout recognition and OCR pipelines with
+ a single model that can take into account both the textual and visual content (layout, charts, ...) of a document.
+
+ ColQwen2 is part of the ColVision model family, which was introduced with ColPali in the following paper:
+ [*ColPali: Efficient Document Retrieval with Vision Language Models*](https://huggingface.co/papers/2407.01449).
+ """
+)
+class ColQwen2ForRetrieval(ColQwen2PreTrainedModel):
+ base_model_prefix = "vlm"
+
+ def __init__(self, config: ColQwen2Config):
+ super().__init__(config)
+ self.config = config
+ self.vocab_size = config.vlm_config.text_config.vocab_size
+
+ self.vlm = AutoModel.from_config(config.vlm_config)
+
+ self.embedding_dim = self.config.embedding_dim
+ self.embedding_proj_layer = nn.Linear(
+ self.config.vlm_config.text_config.hidden_size,
+ self.embedding_dim,
+ )
+
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ labels: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ pixel_values: torch.Tensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ **kwargs,
+ ) -> ColQwen2ForRetrievalOutput:
+ r"""
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ # Handle the custom "pixel_values" input obtained with `ColQwen2Processor` through unpadding
+ if pixel_values is not None and image_grid_thw is not None:
+ # NOTE: image_grid_thw: (batch_size, 3) where image_grid_thw[i] = (num_patches_h, num_patches_w, temporal_patch_size)
+ offsets = image_grid_thw[:, 1] * image_grid_thw[:, 2] # (batch_size,)
+ arange = torch.arange(pixel_values.shape[1], device=offsets.device) # (max_len,)
+ mask = arange.unsqueeze(0) < offsets.unsqueeze(1) # (batch_size, max_len)
+ pixel_values = pixel_values[mask] # (total_valid_patches, channels, height, width)
+
+ 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.return_dict
+
+ # Custom data preparation to fix an issue with the gradient flow when training with multiple GPUs.
+ if inputs_embeds is None:
+ inputs_embeds = self.vlm.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_embeds = self.vlm.visual(pixel_values, grid_thw=image_grid_thw, return_dict=True).pooler_output
+ image_mask = (
+ (input_ids == self.config.vlm_config.image_token_id).unsqueeze(-1).expand_as(inputs_embeds)
+ )
+ image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+ vlm_output = self.vlm(
+ input_ids=None,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ vlm_hidden_states = vlm_output.hidden_states if output_hidden_states else None
+
+ last_hidden_states = vlm_output[0] # (batch_size, sequence_length, hidden_size)
+ proj_dtype = self.embedding_proj_layer.weight.dtype
+ embeddings = self.embedding_proj_layer(last_hidden_states.to(proj_dtype)) # (batch_size, sequence_length, dim)
+
+ # L2 normalization
+ embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True) # (batch_size, sequence_length, dim)
+ if attention_mask is not None:
+ embeddings = embeddings * attention_mask.unsqueeze(-1) # (batch_size, sequence_length, dim)
+
+ return ColQwen2ForRetrievalOutput(
+ embeddings=embeddings,
+ past_key_values=vlm_output.past_key_values,
+ hidden_states=vlm_hidden_states,
+ attentions=vlm_output.attentions,
+ )
+
+
+__all__ = ["ColQwen2ForRetrieval", "ColQwen2PreTrainedModel"]
diff --git a/third_party/transformers/src/transformers/models/colqwen2/modular_colqwen2.py b/third_party/transformers/src/transformers/models/colqwen2/modular_colqwen2.py
new file mode 100644
index 0000000000000000000000000000000000000000..aa7a3f48ca6e6fbb440b4fb42e6271d62e62b494
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/colqwen2/modular_colqwen2.py
@@ -0,0 +1,348 @@
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# 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.
+
+from dataclasses import dataclass
+
+from ...cache_utils import Cache
+from ...feature_extraction_utils import BatchFeature
+from ...image_utils import ImageInput, is_valid_image
+from ...processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack
+from ...tokenization_utils_base import PreTokenizedInput, TextInput
+from ...utils import ModelOutput, auto_docstring, can_return_tuple, is_torch_available, logging
+from ..colpali.modeling_colpali import ColPaliForRetrieval, ColPaliPreTrainedModel
+from ..colpali.processing_colpali import ColPaliProcessor
+from .configuration_colqwen2 import ColQwen2Config
+
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+class ColQwen2ProcessorKwargs(ProcessingKwargs, total=False):
+ _defaults = {
+ "text_kwargs": {
+ "padding": "longest",
+ },
+ "images_kwargs": {
+ "data_format": "channels_first",
+ "do_convert_rgb": True,
+ },
+ "common_kwargs": {"return_tensors": "pt"},
+ }
+
+
+class ColQwen2Processor(ColPaliProcessor):
+ def __init__(
+ self,
+ image_processor=None,
+ tokenizer=None,
+ chat_template=None,
+ visual_prompt_prefix: str | None = None,
+ query_prefix: str | None = None,
+ **kwargs,
+ ):
+ r"""
+ visual_prompt_prefix (`str`, *optional*, defaults to `"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|><|endoftext|>"`):
+ A string that gets tokenized and prepended to the image tokens.
+ query_prefix (`str`, *optional*, defaults to `"Query: "`):
+ A prefix to be used for the query.
+ """
+ ProcessorMixin.__init__(self, image_processor, tokenizer, chat_template=chat_template)
+ self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
+ self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token
+
+ self.visual_prompt_prefix = visual_prompt_prefix or (
+ "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|><|endoftext|>"
+ )
+ self.query_prefix = query_prefix or "Query: "
+
+ def __call__(
+ self,
+ images: ImageInput | None = None,
+ text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None,
+ **kwargs: Unpack[ColQwen2ProcessorKwargs],
+ ) -> BatchFeature:
+ r"""
+ Returns:
+ [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+ - **input_ids** -- List of token ids to be fed to a model.
+ - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+ `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+ `None`).
+ - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+ """
+ output_kwargs = self._merge_kwargs(
+ ColQwen2ProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+ suffix = output_kwargs["text_kwargs"].pop("suffix", None)
+
+ return_token_type_ids = suffix is not None
+
+ if text is None and images is None:
+ raise ValueError("Either text or images must be provided")
+ if text is not None and images is not None:
+ raise ValueError("Only one of text or images can be processed at a time")
+
+ if images is not None:
+ if is_valid_image(images):
+ images = [images]
+ elif isinstance(images, list) and is_valid_image(images[0]):
+ pass
+ elif not (isinstance(images, list) and isinstance(images[0], list) and is_valid_image(images[0][0])):
+ raise ValueError("images must be an image, list of images or list of list of images")
+
+ texts_doc = [self.visual_prompt_prefix] * len(images)
+
+ image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
+ image_grid_thw = image_inputs["image_grid_thw"]
+
+ if image_grid_thw is not None:
+ merge_length = self.image_processor.merge_size**2
+ index = 0
+ for i in range(len(texts_doc)):
+ while self.image_token in texts_doc[i]:
+ texts_doc[i] = texts_doc[i].replace(
+ self.image_token, "<|placeholder|>" * (image_grid_thw[index].prod() // merge_length), 1
+ )
+ index += 1
+ texts_doc[i] = texts_doc[i].replace("<|placeholder|>", self.image_token)
+
+ text_inputs = self.tokenizer(
+ texts_doc,
+ return_token_type_ids=False,
+ **output_kwargs["text_kwargs"],
+ )
+
+ return_data = BatchFeature(data={**text_inputs, **image_inputs})
+
+ # NOTE: The following adjustment ensures correct behavior with DDP on multiple GPUs.
+ offsets = return_data["image_grid_thw"][:, 1] * return_data["image_grid_thw"][:, 2] # (batch_size,)
+
+ # Split the pixel_values tensor into a list of tensors, one per image
+ pixel_values = list(
+ torch.split(return_data["pixel_values"], offsets.tolist())
+ ) # [(num_patches_image_0, pixel_values), ..., (num_patches_image_n, pixel_values)]
+
+ # Pad the list of pixel_value tensors to the same length along the sequence dimension
+ return_data["pixel_values"] = torch.nn.utils.rnn.pad_sequence(
+ pixel_values, batch_first=True
+ ) # (batch_size, max_num_patches, pixel_values)
+
+ if return_token_type_ids:
+ labels = return_data["input_ids"].masked_fill(return_data["token_type_ids"] == 0, -100)
+ return_data.update({"labels": labels})
+
+ return return_data
+
+ elif text is not None:
+ if isinstance(text, str):
+ text = [text]
+ elif not (isinstance(text, list) and isinstance(text[0], str)):
+ raise ValueError("Text must be a string or a list of strings")
+
+ if suffix is None:
+ suffix = self.query_augmentation_token * 10
+
+ texts_query: list[str] = []
+
+ for query in text:
+ augmented_query = self.query_prefix + query + suffix
+ texts_query.append(augmented_query)
+
+ batch_query = self.tokenizer(
+ texts_query,
+ return_token_type_ids=False,
+ **output_kwargs["text_kwargs"],
+ )
+
+ return batch_query
+
+ def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs):
+ """
+ Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.
+ Args:
+ image_sizes (`list[list[int]]`, *optional*):
+ The input sizes formatted as (height, width) per each image.
+ Returns:
+ `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided
+ input modalities, along with other useful data.
+ """
+
+ vision_data = {}
+ if image_sizes is not None:
+ images_kwargs = ColQwen2ProcessorKwargs._defaults.get("images_kwargs", {})
+ images_kwargs.update(kwargs)
+ merge_size = images_kwargs.get("merge_size", None) or self.image_processor.merge_size
+
+ num_image_patches = [
+ self.image_processor.get_number_of_image_patches(*image_size, images_kwargs)
+ for image_size in image_sizes
+ ]
+ num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches]
+ vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
+
+ return MultiModalData(**vision_data)
+
+ @property
+ def model_input_names(self):
+ tokenizer_input_names = self.tokenizer.model_input_names
+ image_processor_input_names = self.image_processor.model_input_names
+
+ # ColQwen doesn't process videos. Make a copy of list when removing
+ # otherwise `self.feature_extractor.model_input_names` is also modified
+ image_processor_input_names = [
+ name for name in image_processor_input_names if name not in ["pixel_values_videos", "video_grid_thw"]
+ ]
+ return tokenizer_input_names + image_processor_input_names
+
+
+class ColQwen2PreTrainedModel(ColPaliPreTrainedModel):
+ pass
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for ColQwen2 embeddings output.
+ """
+)
+class ColQwen2ForRetrievalOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Language modeling loss (for next-token prediction).
+ embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ The embeddings of the model.
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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.
+ """
+
+ loss: torch.FloatTensor | None = None
+ embeddings: torch.Tensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ Following the ColPali approach, ColQwen2 leverages VLMs to construct efficient multi-vector embeddings directly
+ from document images (“screenshots”) for document retrieval. The model is trained to maximize the similarity
+ between these document embeddings and the corresponding query embeddings, using the late interaction method
+ introduced in ColBERT.
+
+ Using ColQwen2 removes the need for potentially complex and brittle layout recognition and OCR pipelines with
+ a single model that can take into account both the textual and visual content (layout, charts, ...) of a document.
+
+ ColQwen2 is part of the ColVision model family, which was introduced with ColPali in the following paper:
+ [*ColPali: Efficient Document Retrieval with Vision Language Models*](https://huggingface.co/papers/2407.01449).
+ """
+)
+class ColQwen2ForRetrieval(ColPaliForRetrieval):
+ def __init__(self, config: ColQwen2Config):
+ super().__init__(config)
+ del self._tied_weights_keys
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ labels: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ pixel_values: torch.Tensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ **kwargs,
+ ) -> ColQwen2ForRetrievalOutput:
+ r"""
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ # Handle the custom "pixel_values" input obtained with `ColQwen2Processor` through unpadding
+ if pixel_values is not None and image_grid_thw is not None:
+ # NOTE: image_grid_thw: (batch_size, 3) where image_grid_thw[i] = (num_patches_h, num_patches_w, temporal_patch_size)
+ offsets = image_grid_thw[:, 1] * image_grid_thw[:, 2] # (batch_size,)
+ arange = torch.arange(pixel_values.shape[1], device=offsets.device) # (max_len,)
+ mask = arange.unsqueeze(0) < offsets.unsqueeze(1) # (batch_size, max_len)
+ pixel_values = pixel_values[mask] # (total_valid_patches, channels, height, width)
+
+ 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.return_dict
+
+ # Custom data preparation to fix an issue with the gradient flow when training with multiple GPUs.
+ if inputs_embeds is None:
+ inputs_embeds = self.vlm.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_embeds = self.vlm.visual(pixel_values, grid_thw=image_grid_thw, return_dict=True).pooler_output
+ image_mask = (
+ (input_ids == self.config.vlm_config.image_token_id).unsqueeze(-1).expand_as(inputs_embeds)
+ )
+ image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+ vlm_output = self.vlm(
+ input_ids=None,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ vlm_hidden_states = vlm_output.hidden_states if output_hidden_states else None
+
+ last_hidden_states = vlm_output[0] # (batch_size, sequence_length, hidden_size)
+ proj_dtype = self.embedding_proj_layer.weight.dtype
+ embeddings = self.embedding_proj_layer(last_hidden_states.to(proj_dtype)) # (batch_size, sequence_length, dim)
+
+ # L2 normalization
+ embeddings = embeddings / embeddings.norm(dim=-1, keepdim=True) # (batch_size, sequence_length, dim)
+ if attention_mask is not None:
+ embeddings = embeddings * attention_mask.unsqueeze(-1) # (batch_size, sequence_length, dim)
+
+ return ColQwen2ForRetrievalOutput(
+ embeddings=embeddings,
+ past_key_values=vlm_output.past_key_values,
+ hidden_states=vlm_hidden_states,
+ attentions=vlm_output.attentions,
+ )
+
+
+__all__ = [
+ "ColQwen2ForRetrieval",
+ "ColQwen2PreTrainedModel",
+ "ColQwen2Processor",
+]
diff --git a/third_party/transformers/src/transformers/models/colqwen2/processing_colqwen2.py b/third_party/transformers/src/transformers/models/colqwen2/processing_colqwen2.py
new file mode 100644
index 0000000000000000000000000000000000000000..48af99206afeb55031853eadd875ab24d0f4cfa2
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/colqwen2/processing_colqwen2.py
@@ -0,0 +1,355 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/colqwen2/modular_colqwen2.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_colqwen2.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# 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.
+
+from typing import Optional, Union
+
+from ...feature_extraction_utils import BatchFeature
+from ...image_utils import ImageInput, is_valid_image
+from ...processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin, Unpack
+from ...tokenization_utils_base import PreTokenizedInput, TextInput
+from ...utils import auto_docstring, is_torch_available
+
+
+if is_torch_available():
+ import torch
+
+
+class ColQwen2ProcessorKwargs(ProcessingKwargs, total=False):
+ _defaults = {
+ "text_kwargs": {
+ "padding": "longest",
+ },
+ "images_kwargs": {
+ "data_format": "channels_first",
+ "do_convert_rgb": True,
+ },
+ "common_kwargs": {"return_tensors": "pt"},
+ }
+
+
+@auto_docstring
+class ColQwen2Processor(ProcessorMixin):
+ def __init__(
+ self,
+ image_processor=None,
+ tokenizer=None,
+ chat_template=None,
+ visual_prompt_prefix: str | None = None,
+ query_prefix: str | None = None,
+ **kwargs,
+ ):
+ r"""
+ visual_prompt_prefix (`str`, *optional*, defaults to `"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|><|endoftext|>"`):
+ A string that gets tokenized and prepended to the image tokens.
+ query_prefix (`str`, *optional*, defaults to `"Query: "`):
+ A prefix to be used for the query.
+ """
+ super().__init__(image_processor, tokenizer, chat_template=chat_template)
+ self.image_token = "<|image_pad|>" if not hasattr(tokenizer, "image_token") else tokenizer.image_token
+ self.video_token = "<|video_pad|>" if not hasattr(tokenizer, "video_token") else tokenizer.video_token
+
+ self.visual_prompt_prefix = visual_prompt_prefix or (
+ "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Describe the image.<|im_end|><|endoftext|>"
+ )
+ self.query_prefix = query_prefix or "Query: "
+
+ @auto_docstring
+ def __call__(
+ self,
+ images: ImageInput | None = None,
+ text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None,
+ **kwargs: Unpack[ColQwen2ProcessorKwargs],
+ ) -> BatchFeature:
+ r"""
+ Returns:
+ [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+ - **input_ids** -- List of token ids to be fed to a model.
+ - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+ `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+ `None`).
+ - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+ """
+ output_kwargs = self._merge_kwargs(
+ ColQwen2ProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+ suffix = output_kwargs["text_kwargs"].pop("suffix", None)
+
+ return_token_type_ids = suffix is not None
+
+ if text is None and images is None:
+ raise ValueError("Either text or images must be provided")
+ if text is not None and images is not None:
+ raise ValueError("Only one of text or images can be processed at a time")
+
+ if images is not None:
+ if is_valid_image(images):
+ images = [images]
+ elif isinstance(images, list) and is_valid_image(images[0]):
+ pass
+ elif not (isinstance(images, list) and isinstance(images[0], list) and is_valid_image(images[0][0])):
+ raise ValueError("images must be an image, list of images or list of list of images")
+
+ texts_doc = [self.visual_prompt_prefix] * len(images)
+
+ image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
+ image_grid_thw = image_inputs["image_grid_thw"]
+
+ if image_grid_thw is not None:
+ merge_length = self.image_processor.merge_size**2
+ index = 0
+ for i in range(len(texts_doc)):
+ while self.image_token in texts_doc[i]:
+ texts_doc[i] = texts_doc[i].replace(
+ self.image_token, "<|placeholder|>" * (image_grid_thw[index].prod() // merge_length), 1
+ )
+ index += 1
+ texts_doc[i] = texts_doc[i].replace("<|placeholder|>", self.image_token)
+
+ text_inputs = self.tokenizer(
+ texts_doc,
+ return_token_type_ids=False,
+ **output_kwargs["text_kwargs"],
+ )
+
+ return_data = BatchFeature(data={**text_inputs, **image_inputs})
+
+ # NOTE: The following adjustment ensures correct behavior with DDP on multiple GPUs.
+ offsets = return_data["image_grid_thw"][:, 1] * return_data["image_grid_thw"][:, 2] # (batch_size,)
+
+ # Split the pixel_values tensor into a list of tensors, one per image
+ pixel_values = list(
+ torch.split(return_data["pixel_values"], offsets.tolist())
+ ) # [(num_patches_image_0, pixel_values), ..., (num_patches_image_n, pixel_values)]
+
+ # Pad the list of pixel_value tensors to the same length along the sequence dimension
+ return_data["pixel_values"] = torch.nn.utils.rnn.pad_sequence(
+ pixel_values, batch_first=True
+ ) # (batch_size, max_num_patches, pixel_values)
+
+ if return_token_type_ids:
+ labels = return_data["input_ids"].masked_fill(return_data["token_type_ids"] == 0, -100)
+ return_data.update({"labels": labels})
+
+ return return_data
+
+ elif text is not None:
+ if isinstance(text, str):
+ text = [text]
+ elif not (isinstance(text, list) and isinstance(text[0], str)):
+ raise ValueError("Text must be a string or a list of strings")
+
+ if suffix is None:
+ suffix = self.query_augmentation_token * 10
+
+ texts_query: list[str] = []
+
+ for query in text:
+ augmented_query = self.query_prefix + query + suffix
+ texts_query.append(augmented_query)
+
+ batch_query = self.tokenizer(
+ texts_query,
+ return_token_type_ids=False,
+ **output_kwargs["text_kwargs"],
+ )
+
+ return batch_query
+
+ def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs):
+ """
+ Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.
+ Args:
+ image_sizes (`list[list[int]]`, *optional*):
+ The input sizes formatted as (height, width) per each image.
+ Returns:
+ `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided
+ input modalities, along with other useful data.
+ """
+
+ vision_data = {}
+ if image_sizes is not None:
+ images_kwargs = ColQwen2ProcessorKwargs._defaults.get("images_kwargs", {})
+ images_kwargs.update(kwargs)
+ merge_size = images_kwargs.get("merge_size", None) or self.image_processor.merge_size
+
+ num_image_patches = [
+ self.image_processor.get_number_of_image_patches(*image_size, images_kwargs)
+ for image_size in image_sizes
+ ]
+ num_image_tokens = [(num_patches // merge_size**2) for num_patches in num_image_patches]
+ vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
+
+ return MultiModalData(**vision_data)
+
+ @property
+ def model_input_names(self):
+ tokenizer_input_names = self.tokenizer.model_input_names
+ image_processor_input_names = self.image_processor.model_input_names
+
+ # ColQwen doesn't process videos. Make a copy of list when removing
+ # otherwise `self.feature_extractor.model_input_names` is also modified
+ image_processor_input_names = [
+ name for name in image_processor_input_names if name not in ["pixel_values_videos", "video_grid_thw"]
+ ]
+ return tokenizer_input_names + image_processor_input_names
+
+ @property
+ def query_augmentation_token(self) -> str:
+ """
+ Return the query augmentation token.
+
+ Query augmentation buffers are used as reasoning buffers during inference.
+ """
+ return self.tokenizer.pad_token
+
+ def process_images(
+ self,
+ images: ImageInput | None = None,
+ **kwargs: Unpack[ColQwen2ProcessorKwargs],
+ ) -> BatchFeature:
+ """
+ Prepare for the model one or several image(s). This method is a wrapper around the `__call__` method of the ColQwen2Processor's
+ [`ColQwen2Processor.__call__`].
+
+ This method forwards the `images` and `kwargs` arguments to the image processor.
+
+ Args:
+ images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `list[PIL.Image.Image]`, `list[np.ndarray]`, `list[torch.Tensor]`):
+ The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+ tensor. In case of a NumPy array/PyTorch tensor, each image should be of shape (C, H, W), where C is a
+ number of channels, H and W are image height and width.
+ return_tensors (`str` or [`~utils.TensorType`], *optional*):
+ If set, will return tensors of a particular framework. Acceptable values are:
+
+ - `'pt'`: Return PyTorch `torch.Tensor` objects.
+ - `'np'`: Return NumPy `np.ndarray` objects.
+
+ Returns:
+ [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+ - **input_ids** -- List of token ids to be fed to a model.
+ - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+ `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+ `None`).
+ - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+ """
+ return self.__call__(images=images, **kwargs)
+
+ def process_queries(
+ self,
+ text: TextInput | list[TextInput],
+ **kwargs: Unpack[ColQwen2ProcessorKwargs],
+ ) -> BatchFeature:
+ """
+ Prepare for the model one or several texts. This method is a wrapper around the `__call__` method of the ColQwen2Processor's
+ [`ColQwen2Processor.__call__`].
+
+ This method forwards the `text` and `kwargs` arguments to the tokenizer.
+
+ Args:
+ text (`str`, `list[str]`, `list[list[str]]`):
+ The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+ (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+ `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+ return_tensors (`str` or [`~utils.TensorType`], *optional*):
+ If set, will return tensors of a particular framework. Acceptable values are:
+
+ - `'pt'`: Return PyTorch `torch.Tensor` objects.
+ - `'np'`: Return NumPy `np.ndarray` objects.
+
+ Returns:
+ [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+ - **input_ids** -- List of token ids to be fed to a model.
+ - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+ `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+ `None`).
+ """
+ return self.__call__(text=text, **kwargs)
+
+ def score_retrieval(
+ self,
+ query_embeddings: Union["torch.Tensor", list["torch.Tensor"]],
+ passage_embeddings: Union["torch.Tensor", list["torch.Tensor"]],
+ batch_size: int = 128,
+ output_dtype: Optional["torch.dtype"] = None,
+ output_device: Union["torch.device", str] = "cpu",
+ ) -> "torch.Tensor":
+ """
+ Compute the late-interaction/MaxSim score (ColBERT-like) for the given multi-vector
+ query embeddings (`qs`) and passage embeddings (`ps`). For ColQwen2, a passage is the
+ image of a document page.
+
+ Because the embedding tensors are multi-vector and can thus have different shapes, they
+ should be fed as:
+ (1) a list of tensors, where the i-th tensor is of shape (sequence_length_i, embedding_dim)
+ (2) a single tensor of shape (n_passages, max_sequence_length, embedding_dim) -> usually
+ obtained by padding the list of tensors.
+
+ Args:
+ query_embeddings (`Union[torch.Tensor, list[torch.Tensor]`): Query embeddings.
+ passage_embeddings (`Union[torch.Tensor, list[torch.Tensor]`): Passage embeddings.
+ batch_size (`int`, *optional*, defaults to 128): Batch size for computing scores.
+ output_dtype (`torch.dtype`, *optional*, defaults to `torch.float32`): The dtype of the output tensor.
+ If `None`, the dtype of the input embeddings is used.
+ output_device (`torch.device` or `str`, *optional*, defaults to "cpu"): The device of the output tensor.
+
+ Returns:
+ `torch.Tensor`: A tensor of shape `(n_queries, n_passages)` containing the scores. The score
+ tensor is saved on the "cpu" device.
+ """
+
+ if len(query_embeddings) == 0:
+ raise ValueError("No queries provided")
+ if len(passage_embeddings) == 0:
+ raise ValueError("No passages provided")
+
+ if query_embeddings[0].device != passage_embeddings[0].device:
+ raise ValueError("Queries and passages must be on the same device")
+
+ if query_embeddings[0].dtype != passage_embeddings[0].dtype:
+ raise ValueError("Queries and passages must have the same dtype")
+
+ if output_dtype is None:
+ output_dtype = query_embeddings[0].dtype
+
+ scores: list[torch.Tensor] = []
+
+ for i in range(0, len(query_embeddings), batch_size):
+ batch_scores: list[torch.Tensor] = []
+ batch_queries = torch.nn.utils.rnn.pad_sequence(
+ query_embeddings[i : i + batch_size], batch_first=True, padding_value=0
+ )
+ for j in range(0, len(passage_embeddings), batch_size):
+ batch_passages = torch.nn.utils.rnn.pad_sequence(
+ passage_embeddings[j : j + batch_size], batch_first=True, padding_value=0
+ )
+ batch_scores.append(
+ torch.einsum("bnd,csd->bcns", batch_queries, batch_passages).max(dim=3)[0].sum(dim=2)
+ )
+ scores.append(torch.cat(batch_scores, dim=1).to(output_dtype).to(output_device))
+
+ return torch.cat(scores, dim=0)
+
+
+__all__ = ["ColQwen2Processor"]
diff --git a/third_party/transformers/src/transformers/models/convnext/__init__.py b/third_party/transformers/src/transformers/models/convnext/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..495f4d40c0f8bed18b58694c5f1d5f35eab82abc
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/convnext/__init__.py
@@ -0,0 +1,29 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_convnext import *
+ from .image_processing_convnext import *
+ from .image_processing_pil_convnext import *
+ from .modeling_convnext import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/convnext/configuration_convnext.py b/third_party/transformers/src/transformers/models/convnext/configuration_convnext.py
new file mode 100644
index 0000000000000000000000000000000000000000..1e1be1995ebcf4e2b05fb8c19c09c4e1d64f72fe
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/convnext/configuration_convnext.py
@@ -0,0 +1,68 @@
+# Copyright 2022 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""ConvNeXT model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...backbone_utils import BackboneConfigMixin
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="facebook/convnext-tiny-224")
+@strict
+class ConvNextConfig(BackboneConfigMixin, PreTrainedConfig):
+ r"""
+ num_stages (`int`, *optional*, defaults to 4):
+ The number of stages in the model.
+
+ Example:
+ ```python
+ >>> from transformers import ConvNextConfig, ConvNextModel
+
+ >>> # Initializing a ConvNext convnext-tiny-224 style configuration
+ >>> configuration = ConvNextConfig()
+
+ >>> # Initializing a model (with random weights) from the convnext-tiny-224 style configuration
+ >>> model = ConvNextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "convnext"
+
+ num_channels: int = 3
+ patch_size: int | list[int] | tuple[int, int] = 4
+ num_stages: int = 4
+ hidden_sizes: list[int] | tuple[int, ...] | None = (96, 192, 384, 768)
+ depths: list[int] | tuple[int, ...] | None = (3, 3, 9, 3)
+ hidden_act: str = "gelu"
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-12
+ layer_scale_init_value: float = 1e-6
+ drop_path_rate: float | int = 0.0
+ image_size: int | list[int] | tuple[int, int] = 224
+ _out_features: list[str] | None = None
+ _out_indices: list[int] | None = None
+
+ def __post_init__(self, **kwargs):
+ self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, len(self.depths) + 1)]
+ self.set_output_features_output_indices(
+ out_indices=kwargs.pop("out_indices", None), out_features=kwargs.pop("out_features", None)
+ )
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["ConvNextConfig"]
diff --git a/third_party/transformers/src/transformers/models/convnext/convert_convnext_to_pytorch.py b/third_party/transformers/src/transformers/models/convnext/convert_convnext_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..caafc1232471d60ac026cba628c284a1d269017e
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/convnext/convert_convnext_to_pytorch.py
@@ -0,0 +1,239 @@
+# Copyright 2022 The HuggingFace Inc. team.
+#
+# 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.
+"""Convert ConvNext checkpoints from the original repository.
+
+URL: https://github.com/facebookresearch/ConvNeXt"""
+
+import argparse
+import json
+from io import BytesIO
+from pathlib import Path
+
+import httpx
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from transformers import ConvNextConfig, ConvNextForImageClassification, ConvNextImageProcessor
+from transformers.utils import logging
+
+
+logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+
+def get_convnext_config(checkpoint_url):
+ config = ConvNextConfig()
+
+ if "tiny" in checkpoint_url:
+ depths = [3, 3, 9, 3]
+ hidden_sizes = [96, 192, 384, 768]
+ if "small" in checkpoint_url:
+ depths = [3, 3, 27, 3]
+ hidden_sizes = [96, 192, 384, 768]
+ if "base" in checkpoint_url:
+ depths = [3, 3, 27, 3]
+ hidden_sizes = [128, 256, 512, 1024]
+ if "large" in checkpoint_url:
+ depths = [3, 3, 27, 3]
+ hidden_sizes = [192, 384, 768, 1536]
+ if "xlarge" in checkpoint_url:
+ depths = [3, 3, 27, 3]
+ hidden_sizes = [256, 512, 1024, 2048]
+
+ if "1k" in checkpoint_url:
+ num_labels = 1000
+ filename = "imagenet-1k-id2label.json"
+ expected_shape = (1, 1000)
+ else:
+ num_labels = 21841
+ filename = "imagenet-22k-id2label.json"
+ expected_shape = (1, 21841)
+
+ repo_id = "huggingface/label-files"
+ config.num_labels = num_labels
+ id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
+ id2label = {int(k): v for k, v in id2label.items()}
+ if "1k" not in checkpoint_url:
+ # this dataset contains 21843 labels but the model only has 21841
+ # we delete the classes as mentioned in https://github.com/google-research/big_transfer/issues/18
+ del id2label[9205]
+ del id2label[15027]
+ config.id2label = id2label
+ config.label2id = {v: k for k, v in id2label.items()}
+ config.hidden_sizes = hidden_sizes
+ config.depths = depths
+
+ return config, expected_shape
+
+
+def rename_key(name):
+ if "downsample_layers.0.0" in name:
+ name = name.replace("downsample_layers.0.0", "embeddings.patch_embeddings")
+ if "downsample_layers.0.1" in name:
+ name = name.replace("downsample_layers.0.1", "embeddings.norm") # we rename to layernorm later on
+ if "downsample_layers.1.0" in name:
+ name = name.replace("downsample_layers.1.0", "stages.1.downsampling_layer.0")
+ if "downsample_layers.1.1" in name:
+ name = name.replace("downsample_layers.1.1", "stages.1.downsampling_layer.1")
+ if "downsample_layers.2.0" in name:
+ name = name.replace("downsample_layers.2.0", "stages.2.downsampling_layer.0")
+ if "downsample_layers.2.1" in name:
+ name = name.replace("downsample_layers.2.1", "stages.2.downsampling_layer.1")
+ if "downsample_layers.3.0" in name:
+ name = name.replace("downsample_layers.3.0", "stages.3.downsampling_layer.0")
+ if "downsample_layers.3.1" in name:
+ name = name.replace("downsample_layers.3.1", "stages.3.downsampling_layer.1")
+ if "stages" in name and "downsampling_layer" not in name:
+ # stages.0.0. for instance should be renamed to stages.0.layers.0.
+ name = name[: len("stages.0")] + ".layers" + name[len("stages.0") :]
+ if "stages" in name:
+ name = name.replace("stages", "encoder.stages")
+ if "norm" in name:
+ name = name.replace("norm", "layernorm")
+ if "gamma" in name:
+ name = name.replace("gamma", "layer_scale_parameter")
+ if "head" in name:
+ name = name.replace("head", "classifier")
+
+ return name
+
+
+# We will verify our results on an image of cute cats
+def prepare_img():
+ url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ with httpx.stream("GET", url) as response:
+ image = Image.open(BytesIO(response.read()))
+ return image
+
+
+@torch.no_grad()
+def convert_convnext_checkpoint(checkpoint_url, pytorch_dump_folder_path):
+ """
+ Copy/paste/tweak model's weights to our ConvNext structure.
+ """
+
+ # define ConvNext configuration based on URL
+ config, expected_shape = get_convnext_config(checkpoint_url)
+ # load original state_dict from URL
+ state_dict = torch.hub.load_state_dict_from_url(checkpoint_url)["model"]
+ # rename keys
+ for key in state_dict.copy():
+ val = state_dict.pop(key)
+ state_dict[rename_key(key)] = val
+ # add prefix to all keys expect classifier head
+ for key in state_dict.copy():
+ val = state_dict.pop(key)
+ if not key.startswith("classifier"):
+ key = "convnext." + key
+ state_dict[key] = val
+
+ # load HuggingFace model
+ model = ConvNextForImageClassification(config)
+ model.load_state_dict(state_dict)
+ model.eval()
+
+ # Check outputs on an image, prepared by ConvNextImageProcessor
+ size = 224 if "224" in checkpoint_url else 384
+ image_processor = ConvNextImageProcessor(size=size)
+ pixel_values = image_processor(images=prepare_img(), return_tensors="pt").pixel_values
+
+ logits = model(pixel_values).logits
+
+ # note: the logits below were obtained without center cropping
+ if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth":
+ expected_logits = torch.tensor([-0.1210, -0.6605, 0.1918])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_small_1k_224_ema.pth":
+ expected_logits = torch.tensor([-0.4473, -0.1847, -0.6365])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_224_ema.pth":
+ expected_logits = torch.tensor([0.4525, 0.7539, 0.0308])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_1k_384.pth":
+ expected_logits = torch.tensor([0.3561, 0.6350, -0.0384])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_224_ema.pth":
+ expected_logits = torch.tensor([0.4174, -0.0989, 0.1489])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_1k_384.pth":
+ expected_logits = torch.tensor([0.2513, -0.1349, -0.1613])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_224.pth":
+ expected_logits = torch.tensor([1.2980, 0.3631, -0.1198])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_224.pth":
+ expected_logits = torch.tensor([1.2963, 0.1227, 0.1723])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_224.pth":
+ expected_logits = torch.tensor([1.7956, 0.8390, 0.2820])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_224.pth":
+ expected_logits = torch.tensor([-0.2822, -0.0502, -0.0878])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_base_22k_1k_384.pth":
+ expected_logits = torch.tensor([-0.5672, -0.0730, -0.4348])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_224.pth":
+ expected_logits = torch.tensor([0.2681, 0.2365, 0.6246])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_large_22k_1k_384.pth":
+ expected_logits = torch.tensor([-0.2642, 0.3931, 0.5116])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_224_ema.pth":
+ expected_logits = torch.tensor([-0.6677, -0.1873, -0.8379])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnext_xlarge_22k_1k_384_ema.pth":
+ expected_logits = torch.tensor([-0.7749, -0.2967, -0.6444])
+ else:
+ raise ValueError(f"Unknown URL: {checkpoint_url}")
+
+ assert torch.allclose(logits[0, :3], expected_logits, atol=1e-3)
+ assert logits.shape == expected_shape
+
+ Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
+ print(f"Saving model to {pytorch_dump_folder_path}")
+ model.save_pretrained(pytorch_dump_folder_path)
+ print(f"Saving image processor to {pytorch_dump_folder_path}")
+ image_processor.save_pretrained(pytorch_dump_folder_path)
+
+ print("Pushing model to the hub...")
+ model_name = "convnext"
+ if "tiny" in checkpoint_url:
+ model_name += "-tiny"
+ elif "small" in checkpoint_url:
+ model_name += "-small"
+ elif "base" in checkpoint_url:
+ model_name += "-base"
+ elif "xlarge" in checkpoint_url:
+ model_name += "-xlarge"
+ elif "large" in checkpoint_url:
+ model_name += "-large"
+ if "224" in checkpoint_url:
+ model_name += "-224"
+ elif "384" in checkpoint_url:
+ model_name += "-384"
+ if "22k" in checkpoint_url and "1k" not in checkpoint_url:
+ model_name += "-22k"
+ if "22k" in checkpoint_url and "1k" in checkpoint_url:
+ model_name += "-22k-1k"
+
+ model.push_to_hub(repo_id=f"nielsr/{model_name}")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ # Required parameters
+ parser.add_argument(
+ "--checkpoint_url",
+ default="https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth",
+ type=str,
+ help="URL of the original ConvNeXT checkpoint you'd like to convert.",
+ )
+ parser.add_argument(
+ "--pytorch_dump_folder_path",
+ default=None,
+ type=str,
+ required=True,
+ help="Path to the output PyTorch model directory.",
+ )
+
+ args = parser.parse_args()
+ convert_convnext_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)
diff --git a/third_party/transformers/src/transformers/models/convnext/image_processing_convnext.py b/third_party/transformers/src/transformers/models/convnext/image_processing_convnext.py
new file mode 100644
index 0000000000000000000000000000000000000000..201376123080194a875dd853be34bb6914f7f615
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/convnext/image_processing_convnext.py
@@ -0,0 +1,145 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Image processor class for ConvNeXT."""
+
+import torch
+from torchvision.transforms.v2 import functional as tvF
+
+from ...image_processing_backends import TorchvisionBackend
+from ...image_processing_utils import BatchFeature
+from ...image_transforms import get_resize_output_image_size, group_images_by_shape, reorder_images
+from ...image_utils import (
+ IMAGENET_STANDARD_MEAN,
+ IMAGENET_STANDARD_STD,
+ ChannelDimension,
+ PILImageResampling,
+ SizeDict,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring
+
+
+class ConvNextImageProcessorKwargs(ImagesKwargs, total=False):
+ r"""
+ crop_pct (`float`, *optional*, defaults to `self.crop_pct`):
+ Percentage of the image to crop. Only has an effect if size < 384.
+ """
+
+ crop_pct: float
+
+
+@auto_docstring
+class ConvNextImageProcessor(TorchvisionBackend):
+ """Torchvision backend for ConvNeXT with custom resize."""
+
+ valid_kwargs = ConvNextImageProcessorKwargs
+
+ resample = PILImageResampling.BICUBIC
+ image_mean = IMAGENET_STANDARD_MEAN
+ image_std = IMAGENET_STANDARD_STD
+ size = {"shortest_edge": 384}
+ default_to_square = False
+ do_resize = True
+ do_rescale = True
+ do_normalize = True
+ crop_pct = 224 / 256
+
+ def __init__(self, **kwargs: Unpack[ConvNextImageProcessorKwargs]):
+ super().__init__(**kwargs)
+
+ def resize(
+ self,
+ image: "torch.Tensor",
+ size: SizeDict,
+ resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+ crop_pct: float = 224 / 256,
+ **kwargs,
+ ) -> "torch.Tensor":
+ """Resize with crop_pct support."""
+ if not size.shortest_edge:
+ raise ValueError(f"Size dictionary must contain 'shortest_edge' key. Got {size.keys()}")
+ shortest_edge = size.shortest_edge
+
+ if shortest_edge < 384:
+ # maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
+ resize_shortest_edge = int(shortest_edge / crop_pct)
+ resize_size = get_resize_output_image_size(
+ image, size=resize_shortest_edge, default_to_square=False, input_data_format=ChannelDimension.FIRST
+ )
+ image = super().resize(
+ image,
+ SizeDict(height=resize_size[0], width=resize_size[1]),
+ resample=resample,
+ **kwargs,
+ )
+ # then crop to (shortest_edge, shortest_edge)
+ return self.center_crop(
+ image,
+ SizeDict(height=shortest_edge, width=shortest_edge),
+ **kwargs,
+ )
+ else:
+ # warping (no cropping) when evaluated at 384 or larger
+ return super().resize(
+ image,
+ SizeDict(height=shortest_edge, width=shortest_edge),
+ resample=resample,
+ **kwargs,
+ )
+
+ def _preprocess(
+ self,
+ images: list["torch.Tensor"],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+ do_center_crop: bool,
+ crop_size: SizeDict,
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ do_pad: bool | None,
+ pad_size: SizeDict | None,
+ disable_grouping: bool | None,
+ return_tensors: str | TensorType | None,
+ crop_pct: float = 224 / 256,
+ **kwargs,
+ ) -> BatchFeature:
+ """Custom preprocessing for ConvNeXT."""
+ grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+ resized_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ if do_resize:
+ stacked_images = self.resize(stacked_images, size, resample, crop_pct)
+ resized_images_grouped[shape] = stacked_images
+ resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+ grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
+ processed_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ if do_center_crop:
+ stacked_images = self.center_crop(stacked_images, crop_size)
+ stacked_images = self.rescale_and_normalize(
+ stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
+ )
+ processed_images_grouped[shape] = stacked_images
+
+ processed_images = reorder_images(processed_images_grouped, grouped_images_index)
+
+ return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)
+
+
+__all__ = ["ConvNextImageProcessor"]
diff --git a/third_party/transformers/src/transformers/models/convnext/image_processing_pil_convnext.py b/third_party/transformers/src/transformers/models/convnext/image_processing_pil_convnext.py
new file mode 100644
index 0000000000000000000000000000000000000000..37c656d00c6ba673b5937fcfdc60bc828c3a126b
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/convnext/image_processing_pil_convnext.py
@@ -0,0 +1,136 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Image processor class for ConvNeXT."""
+
+import numpy as np
+
+from ...image_processing_backends import PilBackend
+from ...image_processing_utils import BatchFeature
+from ...image_transforms import get_resize_output_image_size
+from ...image_utils import (
+ IMAGENET_STANDARD_MEAN,
+ IMAGENET_STANDARD_STD,
+ ChannelDimension,
+ PILImageResampling,
+ SizeDict,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring
+
+
+# Adapted from transformers.models.convnext.image_processing_convnext.ConvNextImageProcessorKwargs
+class ConvNextImageProcessorKwargs(ImagesKwargs, total=False):
+ r"""
+ crop_pct (`float`, *optional*, defaults to `self.crop_pct`):
+ Percentage of the image to crop. Only has an effect if size < 384.
+ """
+
+ crop_pct: float
+
+
+@auto_docstring
+class ConvNextImageProcessorPil(PilBackend):
+ """PIL backend for ConvNeXT with custom resize."""
+
+ valid_kwargs = ConvNextImageProcessorKwargs
+
+ resample = PILImageResampling.BICUBIC
+ image_mean = IMAGENET_STANDARD_MEAN
+ image_std = IMAGENET_STANDARD_STD
+ size = {"shortest_edge": 384}
+ default_to_square = False
+ do_resize = True
+ do_rescale = True
+ do_normalize = True
+ crop_pct = 224 / 256
+
+ def __init__(self, **kwargs: Unpack[ConvNextImageProcessorKwargs]):
+ super().__init__(**kwargs)
+
+ def resize(
+ self,
+ image: np.ndarray,
+ size: SizeDict,
+ resample: "PILImageResampling | None",
+ crop_pct: float = 224 / 256,
+ **kwargs,
+ ) -> np.ndarray:
+ """Resize with crop_pct support."""
+ if not size.shortest_edge:
+ raise ValueError(f"Size dictionary must contain 'shortest_edge' key. Got {size.keys()}")
+ shortest_edge = size.shortest_edge
+
+ if shortest_edge < 384:
+ # maintain same ratio, resizing shortest edge to shortest_edge/crop_pct
+ resize_shortest_edge = int(shortest_edge / crop_pct)
+ resize_size = get_resize_output_image_size(
+ image, size=resize_shortest_edge, default_to_square=False, input_data_format=ChannelDimension.FIRST
+ )
+ image = super().resize(
+ image,
+ size=SizeDict(height=resize_size[0], width=resize_size[1]),
+ resample=resample,
+ **kwargs,
+ )
+ # then crop to (shortest_edge, shortest_edge)
+ return super().center_crop(
+ image,
+ size=SizeDict(height=shortest_edge, width=shortest_edge),
+ **kwargs,
+ )
+ else:
+ # warping (no cropping) when evaluated at 384 or larger
+ return super().resize(
+ image,
+ size=SizeDict(height=shortest_edge, width=shortest_edge),
+ resample=resample,
+ **kwargs,
+ )
+
+ def _preprocess(
+ self,
+ images: list[np.ndarray],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | None",
+ do_center_crop: bool,
+ crop_size: SizeDict,
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ do_pad: bool | None,
+ pad_size: SizeDict | None,
+ return_tensors: str | TensorType | None,
+ crop_pct: float = 224 / 256,
+ **kwargs,
+ ) -> BatchFeature:
+ """Custom preprocessing for ConvNeXT."""
+ processed_images = []
+ for image in images:
+ if do_resize:
+ image = self.resize(image, size, resample, crop_pct)
+ if do_center_crop:
+ image = self.center_crop(image, crop_size)
+ if do_rescale:
+ image = self.rescale(image, rescale_factor)
+ if do_normalize:
+ image = self.normalize(image, image_mean, image_std)
+ processed_images.append(image)
+
+ return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors)
+
+
+__all__ = ["ConvNextImageProcessorPil"]
diff --git a/third_party/transformers/src/transformers/models/convnext/modeling_convnext.py b/third_party/transformers/src/transformers/models/convnext/modeling_convnext.py
new file mode 100644
index 0000000000000000000000000000000000000000..f6f24ab0859746f6031137f98156441f877bc41e
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/convnext/modeling_convnext.py
@@ -0,0 +1,410 @@
+# Copyright 2022 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# 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 ConvNext model."""
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...backbone_utils import BackboneMixin, filter_output_hidden_states
+from ...modeling_outputs import (
+ BackboneOutput,
+ BaseModelOutputWithNoAttention,
+ BaseModelOutputWithPoolingAndNoAttention,
+ ImageClassifierOutputWithNoAttention,
+)
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, logging
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_convnext import ConvNextConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from transformers.models.beit.modeling_beit.drop_path
+def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
+ """
+ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+ """
+ if drop_prob == 0.0 or not training:
+ return input
+ keep_prob = 1 - drop_prob
+ shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
+ random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
+ random_tensor.floor_() # binarize
+ output = input.div(keep_prob) * random_tensor
+ return output
+
+
+# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->ConvNext
+class ConvNextDropPath(nn.Module):
+ """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+ def __init__(self, drop_prob: float | None = None) -> None:
+ super().__init__()
+ self.drop_prob = drop_prob
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ return drop_path(hidden_states, self.drop_prob, self.training)
+
+ def extra_repr(self) -> str:
+ return f"p={self.drop_prob}"
+
+
+class ConvNextLayerNorm(nn.LayerNorm):
+ r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
+ The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height,
+ width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width).
+ """
+
+ def __init__(self, normalized_shape, *, eps=1e-6, data_format="channels_last", **kwargs):
+ super().__init__(normalized_shape, eps=eps, **kwargs)
+ if data_format not in ["channels_last", "channels_first"]:
+ raise NotImplementedError(f"Unsupported data format: {data_format}")
+ self.data_format = data_format
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ """
+ Args:
+ features: Tensor of shape (batch_size, channels, height, width) OR (batch_size, height, width, channels)
+ """
+ if self.data_format == "channels_first":
+ features = features.permute(0, 2, 3, 1)
+ features = super().forward(features)
+ features = features.permute(0, 3, 1, 2)
+ else:
+ features = super().forward(features)
+ return features
+
+
+class ConvNextEmbeddings(nn.Module):
+ """This class is comparable to (and inspired by) the SwinEmbeddings class
+ found in src/transformers/models/swin/modeling_swin.py.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ self.patch_embeddings = nn.Conv2d(
+ config.num_channels, config.hidden_sizes[0], kernel_size=config.patch_size, stride=config.patch_size
+ )
+ self.layernorm = ConvNextLayerNorm(config.hidden_sizes[0], eps=1e-6, data_format="channels_first")
+ self.num_channels = config.num_channels
+
+ def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+ num_channels = pixel_values.shape[1]
+ if num_channels != self.num_channels:
+ raise ValueError(
+ "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+ )
+ embeddings = self.patch_embeddings(pixel_values)
+ embeddings = self.layernorm(embeddings)
+ return embeddings
+
+
+class ConvNextLayer(nn.Module):
+ """This corresponds to the `Block` class in the original implementation.
+
+ There are two equivalent implementations: [DwConv, LayerNorm (channels_first), Conv, GELU,1x1 Conv]; all in (N, C,
+ H, W) (2) [DwConv, Permute to (N, H, W, C), LayerNorm (channels_last), Linear, GELU, Linear]; Permute back
+
+ The authors used (2) as they find it slightly faster in PyTorch.
+
+ Args:
+ config ([`ConvNextConfig`]): Model configuration class.
+ dim (`int`): Number of input channels.
+ drop_path (`float`): Stochastic depth rate. Default: 0.0.
+ """
+
+ def __init__(self, config, dim, drop_path=0):
+ super().__init__()
+ self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim) # depthwise conv
+ self.layernorm = ConvNextLayerNorm(dim, eps=1e-6)
+ self.pwconv1 = nn.Linear(dim, 4 * dim) # pointwise/1x1 convs, implemented with linear layers
+ self.act = ACT2FN[config.hidden_act]
+ self.pwconv2 = nn.Linear(4 * dim, dim)
+ self.layer_scale_parameter = (
+ nn.Parameter(config.layer_scale_init_value * torch.ones(dim), requires_grad=True)
+ if config.layer_scale_init_value > 0
+ else None
+ )
+ self.drop_path = ConvNextDropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ residual = features
+ features = self.dwconv(features)
+ features = features.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C)
+ features = self.layernorm(features)
+ features = self.pwconv1(features)
+ features = self.act(features)
+ features = self.pwconv2(features)
+ if self.layer_scale_parameter is not None:
+ features = self.layer_scale_parameter * features
+ features = features.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W)
+ features = residual + self.drop_path(features)
+ return features
+
+
+class ConvNextStage(nn.Module):
+ """ConvNeXT stage, consisting of an optional downsampling layer + multiple residual blocks.
+
+ Args:
+ config ([`ConvNextConfig`]): Model configuration class.
+ in_channels (`int`): Number of input channels.
+ out_channels (`int`): Number of output channels.
+ depth (`int`): Number of residual blocks.
+ drop_path_rates(`list[float]`): Stochastic depth rates for each layer.
+ """
+
+ def __init__(self, config, in_channels, out_channels, kernel_size=2, stride=2, depth=2, drop_path_rates=None):
+ super().__init__()
+
+ if in_channels != out_channels or stride > 1:
+ self.downsampling_layer = nn.ModuleList(
+ [
+ ConvNextLayerNorm(in_channels, eps=1e-6, data_format="channels_first"),
+ nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride),
+ ]
+ )
+ else:
+ self.downsampling_layer = nn.ModuleList()
+ drop_path_rates = drop_path_rates or [0.0] * depth
+ self.layers = nn.ModuleList(
+ [ConvNextLayer(config, dim=out_channels, drop_path=drop_path_rates[j]) for j in range(depth)]
+ )
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ for layer in self.downsampling_layer:
+ features = layer(features)
+ for layer in self.layers:
+ features = layer(features)
+ return features
+
+
+@auto_docstring
+class ConvNextPreTrainedModel(PreTrainedModel):
+ config: ConvNextConfig
+ base_model_prefix = "convnext"
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+ _no_split_modules = ["ConvNextLayer", "ConvNextStage"]
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ super()._init_weights(module)
+ if isinstance(module, ConvNextLayer):
+ if module.layer_scale_parameter is not None:
+ init.constant_(module.layer_scale_parameter, self.config.layer_scale_init_value)
+
+
+class ConvNextEncoder(ConvNextPreTrainedModel):
+ main_input_name = "hidden_states"
+ _can_record_outputs = {"hidden_states": ConvNextStage}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.stages = nn.ModuleList()
+ drop_path_rates = [
+ x.tolist()
+ for x in torch.linspace(0, config.drop_path_rate, sum(config.depths), device="cpu").split(config.depths)
+ ]
+ prev_chs = config.hidden_sizes[0]
+ for i in range(config.num_stages):
+ out_chs = config.hidden_sizes[i]
+ stage = ConvNextStage(
+ config,
+ in_channels=prev_chs,
+ out_channels=out_chs,
+ stride=2 if i > 0 else 1,
+ depth=config.depths[i],
+ drop_path_rates=drop_path_rates[i],
+ )
+ self.stages.append(stage)
+ prev_chs = out_chs
+
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithNoAttention:
+ for layer_module in self.stages:
+ hidden_states = layer_module(hidden_states)
+
+ return BaseModelOutputWithNoAttention(last_hidden_state=hidden_states)
+
+
+@auto_docstring
+class ConvNextModel(ConvNextPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.config = config
+
+ self.embeddings = ConvNextEmbeddings(config)
+ self.encoder = ConvNextEncoder(config)
+
+ # final layernorm layer
+ self.layernorm = nn.LayerNorm(config.hidden_sizes[-1], eps=config.layer_norm_eps)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self, pixel_values: torch.FloatTensor | None = None, **kwargs: Unpack[TransformersKwargs]
+ ) -> BaseModelOutputWithPoolingAndNoAttention:
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ embedding_output = self.embeddings(pixel_values)
+ encoder_outputs: BaseModelOutputWithNoAttention = self.encoder(embedding_output, **kwargs)
+ last_hidden_state = encoder_outputs.last_hidden_state
+
+ # global average pooling, (N, C, H, W) -> (N, C)
+ pooled_output = self.layernorm(last_hidden_state.mean([-2, -1]))
+
+ return BaseModelOutputWithPoolingAndNoAttention(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ ConvNext Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
+ ImageNet.
+ """
+)
+class ConvNextForImageClassification(ConvNextPreTrainedModel):
+ accepts_loss_kwargs = False
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.num_labels = config.num_labels
+ self.convnext = ConvNextModel(config)
+
+ # Classifier head
+ if config.num_labels > 0:
+ self.classifier = nn.Linear(config.hidden_sizes[-1], config.num_labels)
+ else:
+ self.classifier = nn.Identity()
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self, pixel_values: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, **kwargs
+ ) -> ImageClassifierOutputWithNoAttention:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ outputs: BaseModelOutputWithPoolingAndNoAttention = self.convnext(pixel_values, **kwargs)
+ pooled_output = outputs.pooler_output
+ logits = self.classifier(pooled_output)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(labels=labels, pooled_logits=logits, config=self.config)
+
+ return ImageClassifierOutputWithNoAttention(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ ConvNeXt backbone, to be used with frameworks like DETR and MaskFormer.
+ """
+)
+class ConvNextBackbone(BackboneMixin, ConvNextPreTrainedModel):
+ has_attentions = False
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.embeddings = ConvNextEmbeddings(config)
+ self.encoder = ConvNextEncoder(config)
+ self.num_features = [config.hidden_sizes[0]] + config.hidden_sizes
+
+ # Add layer norms to hidden states of out_features
+ hidden_states_norms = {}
+ for stage, num_channels in zip(self.out_features, self.channels):
+ hidden_states_norms[stage] = ConvNextLayerNorm(num_channels, data_format="channels_first")
+ self.hidden_states_norms = nn.ModuleDict(hidden_states_norms)
+
+ # initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @filter_output_hidden_states
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BackboneOutput:
+ r"""
+ Examples:
+
+ ```python
+ >>> from transformers import AutoImageProcessor, AutoBackbone
+ >>> import torch
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> processor = AutoImageProcessor.from_pretrained("facebook/convnext-tiny-224")
+ >>> model = AutoBackbone.from_pretrained("facebook/convnext-tiny-224")
+
+ >>> inputs = processor(image, return_tensors="pt")
+ >>> outputs = model(**inputs)
+ ```"""
+ kwargs["output_hidden_states"] = True # required to extract layers for the stages
+
+ embedding_output = self.embeddings(pixel_values)
+ encoder_outputs: BaseModelOutputWithNoAttention = self.encoder(embedding_output, **kwargs)
+ hidden_states = encoder_outputs.hidden_states
+
+ feature_maps = []
+ for stage, hidden_state in zip(self.stage_names, hidden_states):
+ if stage in self.out_features:
+ hidden_state = self.hidden_states_norms[stage](hidden_state)
+ feature_maps.append(hidden_state)
+
+ return BackboneOutput(feature_maps=tuple(feature_maps), hidden_states=hidden_states)
+
+
+__all__ = ["ConvNextForImageClassification", "ConvNextModel", "ConvNextPreTrainedModel", "ConvNextBackbone"]
diff --git a/third_party/transformers/src/transformers/models/convnextv2/__init__.py b/third_party/transformers/src/transformers/models/convnextv2/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9e02170eceae73027828d4f184be8e830c02cb8f
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/convnextv2/__init__.py
@@ -0,0 +1,27 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_convnextv2 import *
+ from .modeling_convnextv2 import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/convnextv2/configuration_convnextv2.py b/third_party/transformers/src/transformers/models/convnextv2/configuration_convnextv2.py
new file mode 100644
index 0000000000000000000000000000000000000000..a0c7f9d8346638ac5dece5df7d3a598f6306eca9
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/convnextv2/configuration_convnextv2.py
@@ -0,0 +1,67 @@
+# Copyright 2023 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""ConvNeXTV2 model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...backbone_utils import BackboneConfigMixin
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="facebook/convnextv2-tiny-1k-224")
+@strict
+class ConvNextV2Config(BackboneConfigMixin, PreTrainedConfig):
+ r"""
+ num_stages (`int`, *optional*, defaults to 4):
+ The number of stages in the model.
+
+ Example:
+ ```python
+ >>> from transformers import ConvNeXTV2Config, ConvNextV2Model
+
+ >>> # Initializing a ConvNeXTV2 convnextv2-tiny-1k-224 style configuration
+ >>> configuration = ConvNeXTV2Config()
+
+ >>> # Initializing a model (with random weights) from the convnextv2-tiny-1k-224 style configuration
+ >>> model = ConvNextV2Model(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "convnextv2"
+
+ num_channels: int = 3
+ patch_size: int | list[int] | tuple[int, int] = 4
+ num_stages: int = 4
+ hidden_sizes: list[int] | tuple[int, ...] | None = (96, 192, 384, 768)
+ depths: list[int] | tuple[int, ...] | None = (3, 3, 9, 3)
+ hidden_act: str = "gelu"
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-12
+ drop_path_rate: float | int = 0.0
+ image_size: int | list[int] | tuple[int, int] = 224
+ _out_features: list[str] | None = None
+ _out_indices: list[int] | None = None
+
+ def __post_init__(self, **kwargs):
+ self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, len(self.depths) + 1)]
+ self.set_output_features_output_indices(
+ out_indices=kwargs.pop("out_indices", None), out_features=kwargs.pop("out_features", None)
+ )
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["ConvNextV2Config"]
diff --git a/third_party/transformers/src/transformers/models/convnextv2/convert_convnextv2_to_pytorch.py b/third_party/transformers/src/transformers/models/convnextv2/convert_convnextv2_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..128b4a990c32f2c32c74c9f3c96052c6f16b138c
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/convnextv2/convert_convnextv2_to_pytorch.py
@@ -0,0 +1,287 @@
+# Copyright 2023 The HuggingFace Inc. team.
+#
+# 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.
+"""Convert ConvNeXTV2 checkpoints from the original repository.
+
+URL: https://github.com/facebookresearch/ConvNeXt"""
+
+import argparse
+import json
+import os
+from io import BytesIO
+
+import httpx
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from transformers import ConvNextImageProcessor, ConvNextV2Config, ConvNextV2ForImageClassification
+from transformers.image_utils import PILImageResampling
+from transformers.utils import logging
+
+
+logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+
+def get_convnextv2_config(checkpoint_url):
+ config = ConvNextV2Config()
+
+ if "atto" in checkpoint_url:
+ depths = [2, 2, 6, 2]
+ hidden_sizes = [40, 80, 160, 320]
+ if "femto" in checkpoint_url:
+ depths = [2, 2, 6, 2]
+ hidden_sizes = [48, 96, 192, 384]
+ if "pico" in checkpoint_url:
+ depths = [2, 2, 6, 2]
+ hidden_sizes = [64, 128, 256, 512]
+ if "nano" in checkpoint_url:
+ depths = [2, 2, 8, 2]
+ hidden_sizes = [80, 160, 320, 640]
+ if "tiny" in checkpoint_url:
+ depths = [3, 3, 9, 3]
+ hidden_sizes = [96, 192, 384, 768]
+ if "base" in checkpoint_url:
+ depths = [3, 3, 27, 3]
+ hidden_sizes = [128, 256, 512, 1024]
+ if "large" in checkpoint_url:
+ depths = [3, 3, 27, 3]
+ hidden_sizes = [192, 384, 768, 1536]
+ if "huge" in checkpoint_url:
+ depths = [3, 3, 27, 3]
+ hidden_sizes = [352, 704, 1408, 2816]
+
+ num_labels = 1000
+ filename = "imagenet-1k-id2label.json"
+ expected_shape = (1, 1000)
+
+ repo_id = "huggingface/label-files"
+ config.num_labels = num_labels
+ id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
+ id2label = {int(k): v for k, v in id2label.items()}
+
+ config.id2label = id2label
+ config.label2id = {v: k for k, v in id2label.items()}
+ config.hidden_sizes = hidden_sizes
+ config.depths = depths
+
+ return config, expected_shape
+
+
+def rename_key(name):
+ if "downsample_layers.0.0" in name:
+ name = name.replace("downsample_layers.0.0", "embeddings.patch_embeddings")
+ if "downsample_layers.0.1" in name:
+ name = name.replace("downsample_layers.0.1", "embeddings.norm") # we rename to layernorm later on
+ if "downsample_layers.1.0" in name:
+ name = name.replace("downsample_layers.1.0", "stages.1.downsampling_layer.0")
+ if "downsample_layers.1.1" in name:
+ name = name.replace("downsample_layers.1.1", "stages.1.downsampling_layer.1")
+ if "downsample_layers.2.0" in name:
+ name = name.replace("downsample_layers.2.0", "stages.2.downsampling_layer.0")
+ if "downsample_layers.2.1" in name:
+ name = name.replace("downsample_layers.2.1", "stages.2.downsampling_layer.1")
+ if "downsample_layers.3.0" in name:
+ name = name.replace("downsample_layers.3.0", "stages.3.downsampling_layer.0")
+ if "downsample_layers.3.1" in name:
+ name = name.replace("downsample_layers.3.1", "stages.3.downsampling_layer.1")
+ if "stages" in name and "downsampling_layer" not in name:
+ # stages.0.0. for instance should be renamed to stages.0.layers.0.
+ name = name[: len("stages.0")] + ".layers" + name[len("stages.0") :]
+ if "gamma" in name:
+ name = name.replace("gamma", "weight")
+ if "beta" in name:
+ name = name.replace("beta", "bias")
+ if "stages" in name:
+ name = name.replace("stages", "encoder.stages")
+ if "norm" in name:
+ name = name.replace("norm", "layernorm")
+ if "head" in name:
+ name = name.replace("head", "classifier")
+
+ return name
+
+
+# We will verify our results on an image of cute cats
+def prepare_img():
+ url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ with httpx.stream("GET", url) as response:
+ image = Image.open(BytesIO(response.read()))
+ return image
+
+
+def convert_preprocessor(checkpoint_url):
+ if "224" in checkpoint_url:
+ size = 224
+ crop_pct = 224 / 256
+ elif "384" in checkpoint_url:
+ size = 384
+ crop_pct = None
+ else:
+ size = 512
+ crop_pct = None
+
+ return ConvNextImageProcessor(
+ size=size,
+ crop_pct=crop_pct,
+ image_mean=[0.485, 0.456, 0.406],
+ image_std=[0.229, 0.224, 0.225],
+ resample=PILImageResampling.BICUBIC,
+ )
+
+
+@torch.no_grad()
+def convert_convnextv2_checkpoint(checkpoint_url, pytorch_dump_folder_path, save_model, push_to_hub):
+ """
+ Copy/paste/tweak model's weights to our ConvNeXTV2 structure.
+ """
+ print("Downloading original model from checkpoint...")
+ # define ConvNeXTV2 configuration based on URL
+ config, expected_shape = get_convnextv2_config(checkpoint_url)
+ # load original state_dict from URL
+ state_dict = torch.hub.load_state_dict_from_url(checkpoint_url)["model"]
+
+ print("Converting model parameters...")
+ # rename keys
+ for key in state_dict.copy():
+ val = state_dict.pop(key)
+ state_dict[rename_key(key)] = val
+ # add prefix to all keys expect classifier head
+ for key in state_dict.copy():
+ val = state_dict.pop(key)
+ if not key.startswith("classifier"):
+ key = "convnextv2." + key
+ state_dict[key] = val
+
+ # load HuggingFace model
+ model = ConvNextV2ForImageClassification(config)
+ model.load_state_dict(state_dict)
+ model.eval()
+
+ # Check outputs on an image, prepared by ConvNextImageProcessor
+ preprocessor = convert_preprocessor(checkpoint_url)
+ inputs = preprocessor(images=prepare_img(), return_tensors="pt")
+ logits = model(**inputs).logits
+
+ # note: the logits below were obtained without center cropping
+ if checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt":
+ expected_logits = torch.tensor([-0.3930, 0.1747, -0.5246, 0.4177, 0.4295])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_femto_1k_224_ema.pt":
+ expected_logits = torch.tensor([-0.1727, -0.5341, -0.7818, -0.4745, -0.6566])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_pico_1k_224_ema.pt":
+ expected_logits = torch.tensor([-0.0333, 0.1563, -0.9137, 0.1054, 0.0381])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_nano_1k_224_ema.pt":
+ expected_logits = torch.tensor([-0.1744, -0.1555, -0.0713, 0.0950, -0.1431])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_tiny_1k_224_ema.pt":
+ expected_logits = torch.tensor([0.9996, 0.1966, -0.4386, -0.3472, 0.6661])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_base_1k_224_ema.pt":
+ expected_logits = torch.tensor([-0.2553, -0.6708, -0.1359, 0.2518, -0.2488])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_large_1k_224_ema.pt":
+ expected_logits = torch.tensor([-0.0673, -0.5627, -0.3753, -0.2722, 0.0178])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_huge_1k_224_ema.pt":
+ expected_logits = torch.tensor([-0.6377, -0.7458, -0.2150, 0.1184, -0.0597])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_224_ema.pt":
+ expected_logits = torch.tensor([1.0799, 0.2322, -0.8860, 1.0219, 0.6231])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_nano_22k_384_ema.pt":
+ expected_logits = torch.tensor([0.3766, 0.4917, -1.1426, 0.9942, 0.6024])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_224_ema.pt":
+ expected_logits = torch.tensor([0.4220, -0.6919, -0.4317, -0.2881, -0.6609])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_tiny_22k_384_ema.pt":
+ expected_logits = torch.tensor([0.1082, -0.8286, -0.5095, 0.4681, -0.8085])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_224_ema.pt":
+ expected_logits = torch.tensor([-0.2419, -0.6221, 0.2176, -0.0980, -0.7527])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_base_22k_384_ema.pt":
+ expected_logits = torch.tensor([0.0391, -0.4371, 0.3786, 0.1251, -0.2784])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_224_ema.pt":
+ expected_logits = torch.tensor([-0.0504, 0.5636, -0.1729, -0.6507, -0.3949])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_large_22k_384_ema.pt":
+ expected_logits = torch.tensor([0.3560, 0.9486, 0.3149, -0.2667, -0.5138])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_384_ema.pt":
+ expected_logits = torch.tensor([-0.2469, -0.4550, -0.5853, -0.0810, 0.0309])
+ elif checkpoint_url == "https://dl.fbaipublicfiles.com/convnext/convnextv2/im22k/convnextv2_huge_22k_512_ema.pt":
+ expected_logits = torch.tensor([-0.3090, 0.0802, -0.0682, -0.1979, -0.2826])
+ else:
+ raise ValueError(f"Unknown URL: {checkpoint_url}")
+
+ assert torch.allclose(logits[0, :5], expected_logits, atol=1e-3)
+ assert logits.shape == expected_shape
+ print("Model outputs match the original results!")
+
+ if save_model:
+ print("Saving model to local...")
+ # Create folder to save model
+ if not os.path.isdir(pytorch_dump_folder_path):
+ os.mkdir(pytorch_dump_folder_path)
+
+ model.save_pretrained(pytorch_dump_folder_path)
+ preprocessor.save_pretrained(pytorch_dump_folder_path)
+
+ model_name = "convnextv2"
+ if "atto" in checkpoint_url:
+ model_name += "-atto"
+ if "femto" in checkpoint_url:
+ model_name += "-femto"
+ if "pico" in checkpoint_url:
+ model_name += "-pico"
+ if "nano" in checkpoint_url:
+ model_name += "-nano"
+ elif "tiny" in checkpoint_url:
+ model_name += "-tiny"
+ elif "base" in checkpoint_url:
+ model_name += "-base"
+ elif "large" in checkpoint_url:
+ model_name += "-large"
+ elif "huge" in checkpoint_url:
+ model_name += "-huge"
+ if "22k" in checkpoint_url and "1k" not in checkpoint_url:
+ model_name += "-22k"
+ elif "22k" in checkpoint_url and "1k" in checkpoint_url:
+ model_name += "-22k-1k"
+ elif "1k" in checkpoint_url:
+ model_name += "-1k"
+ if "224" in checkpoint_url:
+ model_name += "-224"
+ elif "384" in checkpoint_url:
+ model_name += "-384"
+ elif "512" in checkpoint_url:
+ model_name += "-512"
+
+ if push_to_hub:
+ print(f"Pushing {model_name} to the hub...")
+ model.push_to_hub(model_name)
+ preprocessor.push_to_hub(model_name)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ # Required parameters
+ parser.add_argument(
+ "--checkpoint_url",
+ default="https://dl.fbaipublicfiles.com/convnext/convnextv2/im1k/convnextv2_atto_1k_224_ema.pt",
+ type=str,
+ help="URL of the original ConvNeXTV2 checkpoint you'd like to convert.",
+ )
+ parser.add_argument(
+ "--pytorch_dump_folder_path",
+ default="model",
+ type=str,
+ help="Path to the output PyTorch model directory.",
+ )
+ parser.add_argument("--save_model", action="store_true", help="Save model to local")
+ parser.add_argument("--push_to_hub", action="store_true", help="Push model and image preprocessor to the hub")
+
+ args = parser.parse_args()
+ convert_convnextv2_checkpoint(
+ args.checkpoint_url, args.pytorch_dump_folder_path, args.save_model, args.push_to_hub
+ )
diff --git a/third_party/transformers/src/transformers/models/convnextv2/modeling_convnextv2.py b/third_party/transformers/src/transformers/models/convnextv2/modeling_convnextv2.py
new file mode 100644
index 0000000000000000000000000000000000000000..06063271febe9a8bcea71d15a596008c23a5bd17
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/convnextv2/modeling_convnextv2.py
@@ -0,0 +1,434 @@
+# Copyright 2023 Meta Platforms, Inc. and The HuggingFace Inc. team. All rights reserved.
+#
+# 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 ConvNextV2 model."""
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...backbone_utils import BackboneMixin, filter_output_hidden_states
+from ...modeling_outputs import (
+ BackboneOutput,
+ BaseModelOutputWithNoAttention,
+ BaseModelOutputWithPoolingAndNoAttention,
+ ImageClassifierOutputWithNoAttention,
+)
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, logging
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_convnextv2 import ConvNextV2Config
+
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from transformers.models.beit.modeling_beit.drop_path
+def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
+ """
+ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+ """
+ if drop_prob == 0.0 or not training:
+ return input
+ keep_prob = 1 - drop_prob
+ shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
+ random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
+ random_tensor.floor_() # binarize
+ output = input.div(keep_prob) * random_tensor
+ return output
+
+
+# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->ConvNextV2
+class ConvNextV2DropPath(nn.Module):
+ """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+ def __init__(self, drop_prob: float | None = None) -> None:
+ super().__init__()
+ self.drop_prob = drop_prob
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ return drop_path(hidden_states, self.drop_prob, self.training)
+
+ def extra_repr(self) -> str:
+ return f"p={self.drop_prob}"
+
+
+class ConvNextV2GRN(nn.Module):
+ """GRN (Global Response Normalization) layer"""
+
+ def __init__(self, dim: int):
+ super().__init__()
+ self.weight = nn.Parameter(torch.zeros(1, 1, 1, dim))
+ self.bias = nn.Parameter(torch.zeros(1, 1, 1, dim))
+
+ def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+ # Compute and normalize global spatial feature maps
+ global_features = torch.linalg.vector_norm(hidden_states, ord=2, dim=(1, 2), keepdim=True)
+ norm_features = global_features / (global_features.mean(dim=-1, keepdim=True) + 1e-6)
+ hidden_states = self.weight * (hidden_states * norm_features) + self.bias + hidden_states
+
+ return hidden_states
+
+
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextLayerNorm with ConvNext->ConvNextV2
+class ConvNextV2LayerNorm(nn.LayerNorm):
+ r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
+ The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height,
+ width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width).
+ """
+
+ def __init__(self, normalized_shape, *, eps=1e-6, data_format="channels_last", **kwargs):
+ super().__init__(normalized_shape, eps=eps, **kwargs)
+ if data_format not in ["channels_last", "channels_first"]:
+ raise NotImplementedError(f"Unsupported data format: {data_format}")
+ self.data_format = data_format
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ """
+ Args:
+ features: Tensor of shape (batch_size, channels, height, width) OR (batch_size, height, width, channels)
+ """
+ if self.data_format == "channels_first":
+ features = features.permute(0, 2, 3, 1)
+ features = super().forward(features)
+ features = features.permute(0, 3, 1, 2)
+ else:
+ features = super().forward(features)
+ return features
+
+
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextEmbeddings with ConvNext->ConvNextV2
+class ConvNextV2Embeddings(nn.Module):
+ """This class is comparable to (and inspired by) the SwinEmbeddings class
+ found in src/transformers/models/swin/modeling_swin.py.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ self.patch_embeddings = nn.Conv2d(
+ config.num_channels, config.hidden_sizes[0], kernel_size=config.patch_size, stride=config.patch_size
+ )
+ self.layernorm = ConvNextV2LayerNorm(config.hidden_sizes[0], eps=1e-6, data_format="channels_first")
+ self.num_channels = config.num_channels
+
+ def forward(self, pixel_values: torch.FloatTensor) -> torch.Tensor:
+ num_channels = pixel_values.shape[1]
+ if num_channels != self.num_channels:
+ raise ValueError(
+ "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+ )
+ embeddings = self.patch_embeddings(pixel_values)
+ embeddings = self.layernorm(embeddings)
+ return embeddings
+
+
+class ConvNextV2Layer(nn.Module):
+ """This corresponds to the `Block` class in the original implementation.
+
+ There are two equivalent implementations: [DwConv, LayerNorm (channels_first), Conv, GELU,1x1 Conv]; all in (N, C,
+ H, W) (2) [DwConv, Permute to (N, H, W, C), LayerNorm (channels_last), Linear, GELU, Linear]; Permute back
+
+ The authors used (2) as they find it slightly faster in PyTorch.
+
+ Args:
+ config ([`ConvNextV2Config`]): Model configuration class.
+ dim (`int`): Number of input channels.
+ drop_path (`float`): Stochastic depth rate. Default: 0.0.
+ """
+
+ def __init__(self, config, dim, drop_path=0):
+ super().__init__()
+ # depthwise conv
+ self.dwconv = nn.Conv2d(dim, dim, kernel_size=7, padding=3, groups=dim)
+ self.layernorm = ConvNextV2LayerNorm(dim, eps=1e-6)
+ # pointwise/1x1 convs, implemented with linear layers
+ self.pwconv1 = nn.Linear(dim, 4 * dim)
+ self.act = ACT2FN[config.hidden_act]
+ self.grn = ConvNextV2GRN(4 * dim)
+ self.pwconv2 = nn.Linear(4 * dim, dim)
+ self.drop_path = ConvNextV2DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ residual = features
+ features = self.dwconv(features)
+ # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
+ features = features.permute(0, 2, 3, 1)
+ features = self.layernorm(features)
+ features = self.pwconv1(features)
+ features = self.act(features)
+ features = self.grn(features)
+ features = self.pwconv2(features)
+ # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
+ features = features.permute(0, 3, 1, 2)
+
+ features = residual + self.drop_path(features)
+ return features
+
+
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextStage with ConvNeXT->ConvNeXTV2, ConvNext->ConvNextV2
+class ConvNextV2Stage(nn.Module):
+ """ConvNeXTV2 stage, consisting of an optional downsampling layer + multiple residual blocks.
+
+ Args:
+ config ([`ConvNextV2Config`]): Model configuration class.
+ in_channels (`int`): Number of input channels.
+ out_channels (`int`): Number of output channels.
+ depth (`int`): Number of residual blocks.
+ drop_path_rates(`list[float]`): Stochastic depth rates for each layer.
+ """
+
+ def __init__(self, config, in_channels, out_channels, kernel_size=2, stride=2, depth=2, drop_path_rates=None):
+ super().__init__()
+
+ if in_channels != out_channels or stride > 1:
+ self.downsampling_layer = nn.ModuleList(
+ [
+ ConvNextV2LayerNorm(in_channels, eps=1e-6, data_format="channels_first"),
+ nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride),
+ ]
+ )
+ else:
+ self.downsampling_layer = nn.ModuleList()
+ drop_path_rates = drop_path_rates or [0.0] * depth
+ self.layers = nn.ModuleList(
+ [ConvNextV2Layer(config, dim=out_channels, drop_path=drop_path_rates[j]) for j in range(depth)]
+ )
+
+ def forward(self, features: torch.Tensor) -> torch.Tensor:
+ for layer in self.downsampling_layer:
+ features = layer(features)
+ for layer in self.layers:
+ features = layer(features)
+ return features
+
+
+@auto_docstring
+class ConvNextV2PreTrainedModel(PreTrainedModel):
+ config: ConvNextV2Config
+ base_model_prefix = "convnextv2"
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+ _no_split_modules = ["ConvNextV2Layer"]
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ super()._init_weights(module)
+ if isinstance(module, ConvNextV2GRN):
+ init.zeros_(module.weight)
+ init.zeros_(module.bias)
+
+
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextEncoder with CONVNEXT->CONVNEXTV2, ConvNext->ConvNextV2
+class ConvNextV2Encoder(ConvNextV2PreTrainedModel):
+ main_input_name = "hidden_states"
+ _can_record_outputs = {"hidden_states": ConvNextV2Stage}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.stages = nn.ModuleList()
+ drop_path_rates = [
+ x.tolist()
+ for x in torch.linspace(0, config.drop_path_rate, sum(config.depths), device="cpu").split(config.depths)
+ ]
+ prev_chs = config.hidden_sizes[0]
+ for i in range(config.num_stages):
+ out_chs = config.hidden_sizes[i]
+ stage = ConvNextV2Stage(
+ config,
+ in_channels=prev_chs,
+ out_channels=out_chs,
+ stride=2 if i > 0 else 1,
+ depth=config.depths[i],
+ drop_path_rates=drop_path_rates[i],
+ )
+ self.stages.append(stage)
+ prev_chs = out_chs
+
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithNoAttention:
+ for layer_module in self.stages:
+ hidden_states = layer_module(hidden_states)
+
+ return BaseModelOutputWithNoAttention(last_hidden_state=hidden_states)
+
+
+@auto_docstring
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextModel with CONVNEXT->CONVNEXTV2, ConvNext->ConvNextV2
+class ConvNextV2Model(ConvNextV2PreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.config = config
+
+ self.embeddings = ConvNextV2Embeddings(config)
+ self.encoder = ConvNextV2Encoder(config)
+
+ # final layernorm layer
+ self.layernorm = nn.LayerNorm(config.hidden_sizes[-1], eps=config.layer_norm_eps)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self, pixel_values: torch.FloatTensor | None = None, **kwargs: Unpack[TransformersKwargs]
+ ) -> BaseModelOutputWithPoolingAndNoAttention:
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ embedding_output = self.embeddings(pixel_values)
+ encoder_outputs: BaseModelOutputWithNoAttention = self.encoder(embedding_output, **kwargs)
+ last_hidden_state = encoder_outputs.last_hidden_state
+
+ # global average pooling, (N, C, H, W) -> (N, C)
+ pooled_output = self.layernorm(last_hidden_state.mean([-2, -1]))
+
+ return BaseModelOutputWithPoolingAndNoAttention(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ ConvNextV2 Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
+ ImageNet.
+ """
+)
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextForImageClassification with CONVNEXT->CONVNEXTV2,ConvNext->ConvNextV2,convnext->convnextv2
+class ConvNextV2ForImageClassification(ConvNextV2PreTrainedModel):
+ accepts_loss_kwargs = False
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.num_labels = config.num_labels
+ self.convnextv2 = ConvNextV2Model(config)
+
+ # Classifier head
+ if config.num_labels > 0:
+ self.classifier = nn.Linear(config.hidden_sizes[-1], config.num_labels)
+ else:
+ self.classifier = nn.Identity()
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self, pixel_values: torch.FloatTensor | None = None, labels: torch.LongTensor | None = None, **kwargs
+ ) -> ImageClassifierOutputWithNoAttention:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ outputs: BaseModelOutputWithPoolingAndNoAttention = self.convnextv2(pixel_values, **kwargs)
+ pooled_output = outputs.pooler_output
+ logits = self.classifier(pooled_output)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(labels=labels, pooled_logits=logits, config=self.config)
+
+ return ImageClassifierOutputWithNoAttention(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ ConvNeXT V2 backbone, to be used with frameworks like DETR and MaskFormer.
+ """
+)
+# Copied from transformers.models.convnext.modeling_convnext.ConvNextBackbone with CONVNEXT->CONVNEXTV2,ConvNext->ConvNextV2,facebook/convnext-tiny-224->facebook/convnextv2-tiny-1k-224
+class ConvNextV2Backbone(BackboneMixin, ConvNextV2PreTrainedModel):
+ has_attentions = False
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.embeddings = ConvNextV2Embeddings(config)
+ self.encoder = ConvNextV2Encoder(config)
+ self.num_features = [config.hidden_sizes[0]] + config.hidden_sizes
+
+ # Add layer norms to hidden states of out_features
+ hidden_states_norms = {}
+ for stage, num_channels in zip(self.out_features, self.channels):
+ hidden_states_norms[stage] = ConvNextV2LayerNorm(num_channels, data_format="channels_first")
+ self.hidden_states_norms = nn.ModuleDict(hidden_states_norms)
+
+ # initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @filter_output_hidden_states
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BackboneOutput:
+ r"""
+ Examples:
+
+ ```python
+ >>> from transformers import AutoImageProcessor, AutoBackbone
+ >>> import torch
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> processor = AutoImageProcessor.from_pretrained("facebook/convnextv2-tiny-1k-224")
+ >>> model = AutoBackbone.from_pretrained("facebook/convnextv2-tiny-1k-224")
+
+ >>> inputs = processor(image, return_tensors="pt")
+ >>> outputs = model(**inputs)
+ ```"""
+ kwargs["output_hidden_states"] = True # required to extract layers for the stages
+
+ embedding_output = self.embeddings(pixel_values)
+ encoder_outputs: BaseModelOutputWithNoAttention = self.encoder(embedding_output, **kwargs)
+ hidden_states = encoder_outputs.hidden_states
+
+ feature_maps = []
+ for stage, hidden_state in zip(self.stage_names, hidden_states):
+ if stage in self.out_features:
+ hidden_state = self.hidden_states_norms[stage](hidden_state)
+ feature_maps.append(hidden_state)
+
+ return BackboneOutput(feature_maps=tuple(feature_maps), hidden_states=hidden_states)
+
+
+__all__ = ["ConvNextV2ForImageClassification", "ConvNextV2Model", "ConvNextV2PreTrainedModel", "ConvNextV2Backbone"]
diff --git a/third_party/transformers/src/transformers/models/cpmant/__init__.py b/third_party/transformers/src/transformers/models/cpmant/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d92eea75693e728f4c1e9afc199d40751ef6af7a
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/cpmant/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_cpmant import *
+ from .modeling_cpmant import *
+ from .tokenization_cpmant import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/cpmant/configuration_cpmant.py b/third_party/transformers/src/transformers/models/cpmant/configuration_cpmant.py
new file mode 100644
index 0000000000000000000000000000000000000000..9f0479982e2d0af47630c507f85d709d4743db35
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/cpmant/configuration_cpmant.py
@@ -0,0 +1,73 @@
+# Copyright 2022 The OpenBMB Team and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""CPMAnt model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="openbmb/cpm-ant-10b")
+@strict
+class CpmAntConfig(PreTrainedConfig):
+ r"""
+ position_bias_num_buckets (`int`, *optional*, defaults to 512):
+ The number of position_bias buckets.
+ position_bias_max_distance (`int`, *optional*, defaults to 2048):
+ The maximum sequence length that this model might ever be used with. Typically set this to something large
+ just in case (e.g., 512 or 1024 or 2048).
+ prompt_types (`int`, *optional*, defaults to 32):
+ The type of prompt.
+ prompt_length (`int`, *optional*, defaults to 32):
+ The length of prompt.
+ segment_types (`int`, *optional*, defaults to 32):
+ The type of segment.
+
+ Example:
+
+ ```python
+ >>> from transformers import CpmAntModel, CpmAntConfig
+
+ >>> # Initializing a CPMAnt cpm-ant-10b style configuration
+ >>> configuration = CpmAntConfig()
+
+ >>> # Initializing a model from the cpm-ant-10b style configuration
+ >>> model = CpmAntModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "cpmant"
+
+ vocab_size: int = 30720
+ hidden_size: int = 4096
+ num_attention_heads: int = 32
+ dim_head: int = 128
+ dim_ff: int = 10240
+ num_hidden_layers: int = 48
+ dropout_p: float | int = 0.0
+ position_bias_num_buckets: int = 512
+ position_bias_max_distance: int = 2048
+ eps: float = 1e-6
+ init_std: float = 1.0
+ prompt_types: int = 32
+ prompt_length: int = 32
+ segment_types: int = 32
+ use_cache: bool = True
+ tie_word_embeddings: bool = True
+
+
+__all__ = ["CpmAntConfig"]
diff --git a/third_party/transformers/src/transformers/models/cpmant/modeling_cpmant.py b/third_party/transformers/src/transformers/models/cpmant/modeling_cpmant.py
new file mode 100644
index 0000000000000000000000000000000000000000..832dc7f7d096a11d9e9a86f8cc3f565cbba343e2
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/cpmant/modeling_cpmant.py
@@ -0,0 +1,779 @@
+# Copyright 2022 The OpenBMB Team and The HuggingFace Inc. team. All rights reserved.
+#
+# 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 CPMAnt"""
+
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_utils import PreTrainedModel
+from ...utils import auto_docstring, logging
+from .configuration_cpmant import CpmAntConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class CpmAntLayerNorm(nn.Module):
+ """
+ We use Root Mean Square (RMS) Layer Normalization, please see https://huggingface.co/papers/1910.07467 for details."
+ """
+
+ def __init__(self, config: CpmAntConfig):
+ super().__init__()
+
+ self.eps = config.eps
+ self.dim_norm = config.hidden_size
+ self.weight = nn.Parameter(torch.empty(config.hidden_size))
+
+ def forward(self, hidden_states: torch.Tensor):
+ """
+ Args:
+ hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
+ """
+ if hidden_states.size(-1) != self.dim_norm:
+ raise AssertionError("hidden_states.size(-1) != self.dim_norm")
+ old_dtype = hidden_states.dtype
+ variance = hidden_states.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
+ hidden_states = (hidden_states * torch.rsqrt(variance + self.eps)).to(old_dtype) * self.weight
+ return hidden_states
+
+
+class CpmAntAttention(nn.Module):
+ def __init__(self, config: CpmAntConfig, layer_idx=None):
+ super().__init__()
+ self.dim_model = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.dim_head = config.dim_head
+ self.layer_idx = layer_idx
+
+ self.project_q = nn.Linear(self.dim_model, self.num_heads * self.dim_head, bias=False)
+ self.project_k = nn.Linear(self.dim_model, self.num_heads * self.dim_head, bias=False)
+ self.project_v = nn.Linear(self.dim_model, self.num_heads * self.dim_head, bias=False)
+
+ self.attention_out = nn.Linear(self.num_heads * self.dim_head, self.dim_model, bias=False)
+
+ self.softmax = torch.nn.Softmax(dim=-1)
+
+ if config.dropout_p is not None:
+ self.dropout = torch.nn.Dropout(p=config.dropout_p)
+ else:
+ self.dropout = None
+
+ def forward(
+ self,
+ hidden_q: torch.Tensor,
+ hidden_kv: torch.Tensor,
+ attention_mask: torch.BoolTensor,
+ position_bias: torch.Tensor,
+ output_attentions: bool | None = False,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ **kwargs,
+ ):
+ """
+ Args:
+ hidden_q (`torch.Tensor`):
+ Input of transformer block(self-attention block). It can be the raw embedding of a batch of sequences.
+ hidden_kv (`torch.Tensor` of shape `(batch, len_k, dim_model)`)):
+ Tensor *key_value* and *query* of shape `(batch, len_k, dim_model)`
+ attention_mask (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
+ Avoid invalid areas to participate in the calculation of self-attention.
+ position_bias (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
+ Provide positional information to self-attention block.
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers.
+ past_key_values (`Cache`, *optional*):
+ Cached past key and value projection states.
+ 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`).
+ """
+ batch_size = hidden_q.size(0)
+ len_q = hidden_q.size(1)
+ len_k = hidden_kv.size(1)
+
+ query = self.project_q(hidden_q)
+ key = self.project_k(hidden_kv)
+ value = self.project_v(hidden_kv)
+
+ query = query.view(batch_size, len_q, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
+ key = key.view(batch_size, len_k, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
+ value = value.view(batch_size, len_k, self.num_heads, self.dim_head).permute(0, 2, 1, 3)
+
+ if past_key_values is not None:
+ key, value = past_key_values.update(key, value, self.layer_idx)
+ len_k = key.size(-2)
+
+ # (batch_size, num_heads, len_q, dim_head) @ (batch_size, num_heads, dim_head, len_k) -> (batch_size, num_heads, len_q, len_k)
+ score = torch.matmul(query, key.transpose(-1, -2)) / math.sqrt(self.dim_head)
+ score = score + position_bias
+
+ score = torch.masked_fill(
+ score,
+ attention_mask.view(batch_size, 1, len_q, len_k) == torch.tensor(False),
+ torch.scalar_tensor(float("-inf"), device=score.device, dtype=score.dtype),
+ )
+ score = self.softmax(score)
+
+ score = torch.masked_fill(
+ score,
+ attention_mask.view(batch_size, 1, len_q, len_k) == torch.tensor(False),
+ torch.scalar_tensor(0, device=score.device, dtype=score.dtype),
+ )
+ if output_attentions:
+ attn_weights = score
+ else:
+ attn_weights = None
+
+ if self.dropout is not None:
+ score = self.dropout(score)
+
+ # (batch_size, num_heads, len_q, len_k) @ (batch_size, num_heads, len_k, dim_head) -> (batch_size, num_heads, len_q, dim_head)
+ score = torch.matmul(score, value)
+
+ score = score.view(batch_size, self.num_heads, len_q, self.dim_head).permute(0, 2, 1, 3)
+ score = score.contiguous().view(batch_size, len_q, self.num_heads * self.dim_head)
+
+ score = self.attention_out(score)
+
+ return score, attn_weights
+
+
+class CpmAntSelfAttentionBlock(nn.Module):
+ def __init__(self, config: CpmAntConfig, layer_idx=None):
+ super().__init__()
+ self.layernorm_before_attention = CpmAntLayerNorm(config)
+ self.self_attention = CpmAntAttention(config, layer_idx=layer_idx)
+ if config.dropout_p:
+ self.dropout = torch.nn.Dropout(config.dropout_p)
+ else:
+ self.dropout = None
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ position_bias: torch.Tensor | None = None,
+ output_attentions: bool | None = False,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ **kwargs,
+ ):
+ """
+ Args:
+ hidden_states (`torch.Tensor` of shape `(batch, len_seq, dim_model)`):
+ Input of transformer block(self-attention block). It can be the raw embedding of a batch of sequences.
+ attention_mask (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
+ Avoid invalid areas to participate in the calculation of self-attention.
+ position_bias (`torch.Tensor` of shape `(batch, len_seq, len_seq)`):
+ Provide positional information to self-attention block.
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers.
+ past_key_values (`Cache`, *optional*):
+ Cached past key and value projection states.
+ 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`).
+ """
+ outputs = self.layernorm_before_attention(hidden_states)
+ outputs, attn_weights = self.self_attention(
+ outputs,
+ outputs,
+ attention_mask,
+ position_bias,
+ output_attentions,
+ past_key_values,
+ use_cache,
+ )
+
+ if self.dropout is not None:
+ outputs = self.dropout(outputs)
+ hidden_states = hidden_states + outputs
+
+ return hidden_states, attn_weights
+
+
+class CpmAntDenseGatedACT(nn.Module):
+ def __init__(self, config: CpmAntConfig):
+ super().__init__()
+ self.w_0 = nn.Linear(config.hidden_size, config.dim_ff, bias=False)
+ self.w_1 = nn.Linear(config.hidden_size, config.dim_ff, bias=False)
+ self.act = torch.nn.GELU()
+
+ def forward(self, hidden_states: torch.Tensor):
+ """Transform an input tensor from one feature space to another via a nonlinear operation
+
+ Args:
+ hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
+ """
+ gate_score = self.act(self.w_0(hidden_states))
+ hidden_states = self.w_1(hidden_states)
+
+ hidden_states = gate_score * hidden_states
+ return hidden_states
+
+
+class CpmAntFeedForward(nn.Module):
+ def __init__(self, config: CpmAntConfig):
+ super().__init__()
+ self.w_in = CpmAntDenseGatedACT(config)
+ if config.dropout_p is not None:
+ self.dropout = torch.nn.Dropout(config.dropout_p)
+ else:
+ self.dropout = None
+
+ self.w_out = nn.Linear(config.dim_ff, config.hidden_size, bias=False)
+
+ def forward(self, hidden_states: torch.Tensor):
+ """
+ Args:
+ hidden_states (`torch.Tensor` of shape `(batch, seq_len, dim_in)`)
+ """
+ hidden_states = self.w_in(hidden_states)
+
+ if self.dropout is not None:
+ hidden_states = self.dropout(hidden_states)
+
+ hidden_states = self.w_out(hidden_states)
+
+ return hidden_states
+
+
+class CpmAntFFNBlock(nn.Module):
+ def __init__(self, config: CpmAntConfig):
+ super().__init__()
+ self.layernorm_before_ffn = CpmAntLayerNorm(config)
+ self.ffn = CpmAntFeedForward(config)
+ if config.dropout_p:
+ self.dropout = torch.nn.Dropout(config.dropout_p)
+ else:
+ self.dropout = None
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ ):
+ """
+ Args:
+ hidden_states (`torch.Tensor` of shape `(batch, len_seq, dim_model)`):
+ Hidden states before feed forward layer.
+ """
+ ln_outputs = self.layernorm_before_ffn(hidden_states)
+ outputs = self.ffn(ln_outputs)
+ if self.dropout is not None:
+ outputs = self.dropout(outputs)
+ hidden_states = hidden_states + outputs
+ return hidden_states
+
+
+class CpmAntTransformerBlock(nn.Module):
+ def __init__(self, config: CpmAntConfig, layer_idx=None):
+ super().__init__()
+ self.self_att = CpmAntSelfAttentionBlock(config, layer_idx=layer_idx)
+ self.ffn = CpmAntFFNBlock(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ position_bias: torch.Tensor | None = None,
+ output_attentions: bool | None = False,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ **kwargs,
+ ):
+ """
+ Args:
+ hidden_states (`torch.Tensor`):
+ Input to the layer of shape `(batch, seq_len, dim_model)`
+ attention_mask (`torch.Tensor`):
+ Avoid invalid areas to participate in the calculation of shape `(batch, seq_len, seq_len)`
+ position_bias (`torch.Tensor`):
+ Provides position information to attention mechanism of shape `(num_heads, seq_len, seq_len)`
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers.
+ past_key_values (`Cache`, *optional*):
+ Cached past key and value projection states
+ 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`).
+ """
+ hidden_states, attn_weights = self.self_att(
+ hidden_states,
+ attention_mask=attention_mask,
+ position_bias=position_bias,
+ output_attentions=output_attentions,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ )
+
+ hidden_states = self.ffn(hidden_states)
+ return hidden_states, attn_weights
+
+
+class CpmAntEncoder(nn.Module):
+ def __init__(self, config: CpmAntConfig):
+ super().__init__()
+ self.num_layers = config.num_hidden_layers
+ self.layers = nn.ModuleList([CpmAntTransformerBlock(config, layer_idx=i) for i in range(self.num_layers)])
+
+ self.output_layernorm = CpmAntLayerNorm(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ position_bias: torch.Tensor,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ **kwargs,
+ ):
+ """
+ Args:
+ hidden_states (`torch.Tensor`):
+ Input to the layer of shape `(batch, seq_len, dim_model)`
+ attention_mask (`torch.Tensor`):
+ Avoid invalid areas to participate in the calculation of shape `(batch, seq_len, seq_len)`
+ position_bias (`torch.Tensor`):
+ Provides position information to attention mechanism of shape `(num_heads, seq_len, seq_len)`
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers.
+ output_hidden_states (`bool`, *optional*):
+ Whether or not to return the hidden states of all layers.
+ past_key_values (`Cache`, *optional*):
+ Cached past key and value projection states
+ 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`).
+ """
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attns = () if output_attentions else None
+
+ for i, layer in enumerate(self.layers):
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+ layer_outputs = layer(
+ hidden_states,
+ attention_mask,
+ position_bias,
+ output_attentions=output_attentions,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ )
+ hidden_states, attn_weights = layer_outputs
+ if output_attentions:
+ all_self_attns += (attn_weights,)
+
+ hidden_states = self.output_layernorm(hidden_states)
+
+ if output_hidden_states:
+ all_hidden_states += (hidden_states,)
+
+ return hidden_states, all_hidden_states, all_self_attns
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->CPMAnt
+class CpmAntIntermediate(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+class CpmAntSegmentPositionEmbedding(nn.Module):
+ def __init__(self, config: CpmAntConfig):
+ super().__init__()
+
+ self.num_heads = config.num_attention_heads
+ self.num_buckets = config.position_bias_num_buckets
+ self.max_distance = config.position_bias_max_distance
+ self.num_segments = config.segment_types
+
+ self.relative_attention_bias = nn.Parameter(
+ torch.empty(
+ config.segment_types * config.segment_types + config.position_bias_num_buckets,
+ config.num_attention_heads,
+ )
+ )
+
+ def forward(
+ self,
+ key_pos: torch.Tensor,
+ query_pos: torch.Tensor,
+ key_segment: torch.Tensor,
+ query_segment: torch.Tensor,
+ ):
+ with torch.no_grad():
+ batch = key_pos.size(0)
+ keylen = key_pos.size(1)
+ querylen = query_pos.size(1)
+
+ if key_pos.size(0) != query_pos.size(0):
+ raise AssertionError(
+ f"key_pos.size(0) should be equal to query_pos.size(0), but got {key_pos.size(0)} and {query_pos.size(0)}!"
+ )
+ if keylen != key_segment.size(1) or querylen != query_segment.size(1):
+ raise AssertionError(
+ f"keylen should be equal to key_segment.size(1), but got {keylen} and {key_segment.size(1)}!"
+ )
+ if querylen != query_segment.size(1):
+ raise AssertionError(
+ f"querylen should be equal to query_segment.size(1), but got {querylen} and {query_segment.size(1)}!"
+ )
+
+ key_pos = key_pos.view(batch, -1, keylen)
+ query_pos = query_pos.view(batch, querylen, -1)
+ key_segment = key_segment.view(batch, -1, keylen)
+ query_segment = query_segment.view(batch, querylen, -1)
+
+ relative_position_bucket = self._segment_relative_position_bucket(query_segment, key_segment)
+ relative_position_bucket = relative_position_bucket + self.num_buckets
+
+ # (batch, len_q, len_k)
+ absolute_position_bucket = self._position_bucket(
+ torch.arange(keylen, dtype=torch.int32, device=relative_position_bucket.device)[None, :]
+ - torch.arange(querylen, dtype=torch.int32, device=relative_position_bucket.device)[:, None],
+ num_buckets=self.num_buckets,
+ max_distance=self.max_distance,
+ )
+ relative_position_bucket = torch.where(
+ (key_segment == query_segment),
+ absolute_position_bucket[None, :, :],
+ relative_position_bucket,
+ )
+
+ # (batch, len_q, len_k, num_heads)
+ embeds = F.embedding(relative_position_bucket, self.relative_attention_bias)
+ # (batch, num_heads, len_q, len_k)
+ embeds = embeds.permute(0, 3, 1, 2).contiguous()
+ return embeds
+
+ def _segment_relative_position_bucket(self, query_segment, key_segment):
+ return query_segment * self.num_segments + key_segment
+
+ def _position_bucket(self, relative_position, num_buckets=32, max_distance=128):
+ relative_buckets = 0
+ # always bidirectional in CPMAnt
+ num_buckets //= 2
+ relative_buckets = (relative_position > 0).to(torch.int32) * num_buckets
+ relative_position = torch.abs(relative_position)
+ max_exact = num_buckets // 2
+ is_small = relative_position < max_exact
+ relative_position_if_large = max_exact + (
+ torch.log(relative_position.float() / max_exact)
+ / math.log(max_distance / max_exact)
+ * (num_buckets - max_exact)
+ ).to(torch.int32)
+ relative_position_if_large = torch.min(
+ relative_position_if_large,
+ torch.full_like(relative_position_if_large, num_buckets - 1),
+ )
+ relative_buckets += torch.where(is_small, relative_position.to(torch.int32), relative_position_if_large)
+ return relative_buckets
+
+
+# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->CPMAnt
+class CpmAntOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+@auto_docstring
+class CpmAntPreTrainedModel(PreTrainedModel):
+ config: CpmAntConfig
+ base_model_prefix = "cpmant"
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ super()._init_weights(module)
+ if isinstance(module, CpmAntLayerNorm):
+ init.ones_(module.weight)
+ elif isinstance(module, CpmAntSegmentPositionEmbedding):
+ init.normal_(module.relative_attention_bias, mean=0.0, std=self.config.init_std)
+
+
+@auto_docstring
+class CpmAntModel(CpmAntPreTrainedModel):
+ def __init__(self, config: CpmAntConfig):
+ super().__init__(config)
+ self.encoder = CpmAntEncoder(config)
+ self.segment_embedding = nn.Embedding(config.segment_types, config.hidden_size)
+ self.input_embedding = nn.Embedding(
+ config.vocab_size + config.prompt_types * config.prompt_length, config.hidden_size
+ )
+ self.position_bias = CpmAntSegmentPositionEmbedding(config)
+ self.prompt_length = config.prompt_length
+ self.vocab_size = config.vocab_size
+
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.input_embedding
+
+ def set_input_embeddings(self, embeddings, **kwargs):
+ self.input_embedding = embeddings
+
+ def _prepare_attention_mask(self, input_ids, span, context, length):
+ batch = input_ids.size(0)
+ seqlen = input_ids.size(1)
+ device = input_ids.device
+ directional_mask_2d = torch.arange(seqlen, device=device) <= torch.arange(seqlen, device=device).view(-1, 1)
+ attention_mask = context[:, None, :] | (
+ context[:, :, None].logical_not() & directional_mask_2d.view(1, seqlen, seqlen)
+ )
+ attention_mask = attention_mask & (span[:, None, :] == span[:, :, None])
+ # mask for left padding
+ mask_1d = (
+ torch.tensor(list(range(seqlen - self.prompt_length))[::-1], device=device)[None, :].repeat(batch, 1)
+ < length[:, None]
+ )
+ mask_1d = torch.cat((torch.ones(batch, self.prompt_length, device=device).bool(), mask_1d), dim=1)
+ attention_mask = mask_1d.view(batch, seqlen, 1) & mask_1d.view(batch, 1, seqlen) & attention_mask
+ return attention_mask
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor] | BaseModelOutputWithPast:
+ r"""
+ input_ids (`torch.Tensor` of shape `(batch_size, seq_len)`):
+ Indices of input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`CPMAntTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ """
+ 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.return_dict
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+
+ # add prompts ahead
+ if input_ids.dtype != torch.int32:
+ input_ids = input_ids.to(torch.int32)
+ dtype, device = input_ids.dtype, input_ids.device
+ segment = torch.where(input_ids != 0, 2, 0).to(dtype=dtype, device=device)
+ length = (segment != 0).sum(-1).to(dtype=dtype, device=device)
+ input_ids = torch.cat(
+ (
+ torch.arange(
+ self.prompt_length * 2 + self.vocab_size,
+ self.prompt_length * 3 + self.vocab_size,
+ dtype=dtype,
+ device=device,
+ ).repeat(input_ids.size(0), 1),
+ input_ids,
+ ),
+ dim=1,
+ )
+ batch, seq_length = input_ids.size()
+ segment = torch.cat((torch.zeros(batch, self.prompt_length, dtype=dtype, device=device), segment), dim=1)
+ context = torch.full((batch, seq_length), 1, dtype=dtype, device=device)
+ position = torch.arange(seq_length, dtype=dtype, device=device).repeat(batch, 1)
+ span = torch.full((batch, seq_length), 0, dtype=dtype, device=device)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+ input_ids = input_ids.contiguous()
+ hidden_states = self.input_embedding(input_ids)
+ segment_states = self.segment_embedding(segment)
+ if past_length != 0:
+ segment_states = segment_states[:, -1:, :]
+
+ hidden_states = hidden_states + segment_states
+
+ attention_mask = self._prepare_attention_mask(input_ids, span, context, length)
+ position_bias = self.position_bias(position, position, segment, segment)
+
+ attention_mask = attention_mask[:, past_length:, :]
+ position_bias = position_bias[:, :, past_length:, :]
+ hidden_states = hidden_states[:, past_length:, :]
+
+ hidden_states, all_hidden_states, all_attentions = self.encoder(
+ hidden_states,
+ attention_mask,
+ position_bias,
+ output_attentions,
+ output_hidden_states,
+ past_key_values,
+ use_cache,
+ )
+
+ if past_length == 0:
+ hidden_states = hidden_states[:, self.prompt_length :, :]
+ # drop the prompt
+ if all_attentions is not None:
+ new_attentions = ()
+ for attention in all_attentions:
+ new_attentions += (attention[:, :, self.prompt_length :, self.prompt_length :],)
+ all_attentions = new_attentions
+ if all_hidden_states is not None:
+ new_hidden_states = ()
+ for hidden_state in all_hidden_states:
+ new_hidden_states += (hidden_state[:, self.prompt_length :, :],)
+ all_hidden_states = new_hidden_states
+
+ if not return_dict:
+ return tuple(
+ v for v in [hidden_states, past_key_values, all_hidden_states, all_attentions] if v is not None
+ )
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ hidden_states=all_hidden_states,
+ attentions=all_attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The CPMAnt Model with a language modeling head on top (linear layer with weights tied to the input embeddings).
+ """
+)
+class CpmAntForCausalLM(CpmAntPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "cpmant.input_embedding.weight"}
+
+ def __init__(self, config: CpmAntConfig):
+ super().__init__(config)
+ self.cpmant = CpmAntModel(config)
+
+ # lm_head.weight is tied to cpmant.input_embedding.weight
+ self.lm_head = nn.Linear(
+ config.hidden_size, config.vocab_size + config.prompt_types * config.prompt_length, bias=False
+ )
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ labels: torch.Tensor | None = None,
+ return_dict: bool | None = None,
+ attention_mask: torch.Tensor | None = None, # dummy parameter for text-generation pipeline
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs,
+ ) -> tuple | CausalLMOutputWithPast:
+ r"""
+ input_ids (`torch.Tensor` of shape `(batch_size, seq_len)`):
+ Indices of input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`CPMAntTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss.
+
+ Example:
+
+ Text Generation with CpmAntForCausalLM.
+ ```python
+ >>> from transformers import CPMAntTokenizer, CpmAntForCausalLM
+
+ >>> texts = "今天天气不错，"
+ >>> model = CpmAntForCausalLM.from_pretrained("openbmb/cpm-ant-10b")
+ >>> tokenizer = CPMAntTokenizer.from_pretrained("openbmb/cpm-ant-10b")
+ >>> input_ids = tokenizer(texts, return_tensors="pt")
+ >>> outputs = model.generate(**input_ids)
+ >>> output_texts = tokenizer.batch_decode(outputs)
+ >>> print(output_texts)
+ ['今天天气不错，阳光明媚，我和妈妈一起去超市买东西。\n在超市里，我看到了一个很好玩的玩具，它的名字叫“机器人”。它有一个圆圆的脑袋，两只圆圆的眼睛，还有一个圆圆的']
+ ```
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ model_output = self.cpmant(
+ input_ids,
+ output_attentions,
+ output_hidden_states,
+ past_key_values,
+ use_cache,
+ return_dict,
+ )
+ hidden_states = model_output.last_hidden_state if return_dict else model_output[0]
+ # Only compute necessary logits
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss_func = CrossEntropyLoss()
+ loss = loss_func(logits.view(-1, logits.size(-1)), labels.view(-1))
+
+ if not return_dict:
+ output = (logits,) + model_output[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=model_output.past_key_values,
+ hidden_states=model_output.hidden_states,
+ attentions=model_output.attentions,
+ )
+
+ def get_input_embeddings(self):
+ return self.cpmant.input_embedding
+
+ def set_input_embeddings(self, embeddings):
+ self.cpmant.input_embedding = embeddings
+
+
+__all__ = ["CpmAntForCausalLM", "CpmAntModel", "CpmAntPreTrainedModel"]
diff --git a/third_party/transformers/src/transformers/models/cpmant/tokenization_cpmant.py b/third_party/transformers/src/transformers/models/cpmant/tokenization_cpmant.py
new file mode 100644
index 0000000000000000000000000000000000000000..4404a53fcee67a9006082988ec50eadc4229f2e9
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/cpmant/tokenization_cpmant.py
@@ -0,0 +1,232 @@
+# Copyright 2022 The OpenBMB Team and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Tokenization classes for CPMAnt."""
+
+import collections
+import os
+
+from transformers.utils import is_rjieba_available, requires_backends
+
+
+if is_rjieba_available():
+ import rjieba
+
+from ...tokenization_python import PreTrainedTokenizer
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
+
+
+def load_vocab(vocab_file):
+ """Loads a vocabulary file into a dictionary."""
+ vocab = collections.OrderedDict()
+ with open(vocab_file, "r", encoding="utf-8") as reader:
+ tokens = reader.readlines()
+ for index, token in enumerate(tokens):
+ token = token.rstrip("\n")
+ vocab[token] = index
+ return vocab
+
+
+class WordpieceTokenizer:
+ def __init__(self, vocab, unk_token="", max_input_chars_per_word=200):
+ self.vocab = vocab
+ self.unk_token = unk_token
+ self.max_input_chars_per_word = max_input_chars_per_word
+
+ def tokenize(self, token):
+ chars = list(token)
+ if len(chars) > self.max_input_chars_per_word:
+ return [self.unk_token]
+
+ start = 0
+ sub_tokens = []
+ while start < len(chars):
+ end = len(chars)
+ cur_substr = None
+ while start < end:
+ substr = "".join(chars[start:end])
+ if substr in self.vocab:
+ cur_substr = substr
+ break
+ end -= 1
+ if cur_substr is None:
+ sub_tokens.append(self.unk_token)
+ start += 1
+ else:
+ sub_tokens.append(cur_substr)
+ start = end
+
+ return sub_tokens
+
+
+class CpmAntTokenizer(PreTrainedTokenizer):
+ """
+ Construct a CPMAnt tokenizer. Based on byte-level Byte-Pair-Encoding.
+
+ Args:
+ vocab_file (`str`):
+ Path to the vocabulary file.
+ bod_token (`str`, *optional*, defaults to `""`):
+ The beginning of document token.
+ eod_token (`str`, *optional*, defaults to `""`):
+ The end of document token.
+ bos_token (`str`, *optional*, defaults to `""`):
+ The beginning of sequence token.
+ eos_token (`str`, *optional*, defaults to `""`):
+ The end of sequence token.
+ pad_token (`str`, *optional*, defaults to `""`):
+ The token used for padding.
+ unk_token (`str`, *optional*, defaults to `""`):
+ The unknown token.
+ line_token (`str`, *optional*, defaults to `""`):
+ The line token.
+ space_token (`str`, *optional*, defaults to `""`):
+ The space token.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ model_input_names = ["input_ids", "attention_mask"]
+ add_prefix_space = False
+
+ def __init__(
+ self,
+ vocab_file,
+ bod_token="",
+ eod_token="",
+ bos_token="",
+ eos_token="",
+ pad_token="",
+ unk_token="",
+ line_token="",
+ space_token="",
+ padding_side="left",
+ **kwargs,
+ ):
+ requires_backends(self, ["rjieba"])
+ self.bod_token = bod_token
+ self.eod_token = eod_token
+ self.encoder = load_vocab(vocab_file)
+ self.encoder[" "] = self.encoder[space_token]
+ self.encoder["\n"] = self.encoder[line_token]
+
+ del self.encoder[space_token]
+ del self.encoder[line_token]
+
+ self.encoder = collections.OrderedDict(sorted(self.encoder.items(), key=lambda x: x[1]))
+ self.decoder = {v: k for k, v in self.encoder.items()}
+
+ self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.encoder, unk_token=unk_token)
+
+ super().__init__(
+ bod_token=bod_token,
+ eod_token=eod_token,
+ bos_token=bos_token,
+ eos_token=eos_token,
+ pad_token=pad_token,
+ unk_token=unk_token,
+ line_token=line_token,
+ space_token=space_token,
+ padding_side=padding_side,
+ token_type_ids_pattern="all_zeros",
+ token_type_ids_include_special_tokens=True,
+ special_tokens_pattern="bos",
+ **kwargs,
+ )
+ for special_token in [space_token, line_token]:
+ token_id = self.added_tokens_encoder.pop(special_token, None)
+ if token_id is not None:
+ self._added_tokens_decoder.pop(token_id, None)
+ self._update_total_vocab_size()
+
+ @property
+ def bod_token_id(self):
+ return self.encoder[self.bod_token]
+
+ @property
+ def eod_token_id(self):
+ return self.encoder[self.eod_token]
+
+ @property
+ def newline_id(self):
+ return self.encoder["\n"]
+
+ @property
+ def vocab_size(self) -> int:
+ return len(self.encoder)
+
+ def get_vocab(self):
+ return dict(self.encoder, **self.added_tokens_encoder)
+
+ def _tokenize(self, text):
+ """Tokenize a string."""
+ output_tokens = []
+ for x in rjieba.cut(text, False):
+ output_tokens.extend(self.wordpiece_tokenizer.tokenize(x))
+ return output_tokens
+
+ def _decode(self, token_ids, **kwargs):
+ """Decode ids into a string."""
+ token_ids = [i for i in token_ids if i >= 0]
+ token_ids = [
+ x for x in token_ids if x != self.pad_token_id and x != self.eos_token_id and x != self.bos_token_id
+ ]
+ return super()._decode(token_ids, **kwargs)
+
+ def check(self, token):
+ return token in self.encoder
+
+ def convert_tokens_to_string(self, tokens: list[str]) -> str:
+ return "".join(tokens)
+
+ def _convert_token_to_id(self, token):
+ """Converts a token (str) in an id using the vocab."""
+ return self.encoder.get(token, self.encoder.get(self.unk_token))
+
+ def _convert_id_to_token(self, index):
+ """Converts an index (integer) in a token (str) using the vocab."""
+ return self.decoder.get(index, self.unk_token)
+
+ def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None) -> tuple[str]:
+ if os.path.isdir(save_directory):
+ vocab_file = os.path.join(
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
+ )
+ else:
+ vocab_file = (filename_prefix + "-" if filename_prefix else "") + save_directory
+ index = 0
+ if " " in self.encoder:
+ self.encoder[""] = self.encoder[" "]
+ del self.encoder[" "]
+ if "\n" in self.encoder:
+ self.encoder[""] = self.encoder["\n"]
+ del self.encoder["\n"]
+ self.encoder = collections.OrderedDict(sorted(self.encoder.items(), key=lambda x: x[1]))
+ with open(vocab_file, "w", encoding="utf-8") as writer:
+ for token, token_index in self.encoder.items():
+ if index != token_index:
+ logger.warning(
+ f"Saving vocabulary to {vocab_file}: vocabulary indices are not consecutive."
+ " Please check that the vocabulary is not corrupted!"
+ )
+ index = token_index
+ writer.write(token + "\n")
+ index += 1
+ return (vocab_file,)
+
+
+__all__ = ["CpmAntTokenizer"]
diff --git a/third_party/transformers/src/transformers/models/ctrl/__init__.py b/third_party/transformers/src/transformers/models/ctrl/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..93f27ba0710eecb71d1a4ba6b7c306c7ef7d13e8
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/ctrl/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_ctrl import *
+ from .modeling_ctrl import *
+ from .tokenization_ctrl import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/ctrl/configuration_ctrl.py b/third_party/transformers/src/transformers/models/ctrl/configuration_ctrl.py
new file mode 100644
index 0000000000000000000000000000000000000000..8340e6df32bcbf37db7bad05c0d3bbf7375b215a
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/ctrl/configuration_ctrl.py
@@ -0,0 +1,70 @@
+# Copyright 2018 Salesforce and HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+# 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.
+"""Salesforce CTRL configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="Salesforce/ctrl")
+@strict
+class CTRLConfig(PreTrainedConfig):
+ r"""
+ dff (`int`, *optional*, defaults to 8192):
+ Dimensionality of the inner dimension of the feed forward networks (FFN).
+
+ Examples:
+
+ ```python
+ >>> from transformers import CTRLConfig, CTRLModel
+
+ >>> # Initializing a CTRL configuration
+ >>> configuration = CTRLConfig()
+
+ >>> # Initializing a model (with random weights) from the configuration
+ >>> model = CTRLModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "ctrl"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ attribute_map = {
+ "max_position_embeddings": "n_positions",
+ "hidden_size": "n_embd",
+ "num_attention_heads": "n_head",
+ "num_hidden_layers": "n_layer",
+ }
+
+ vocab_size: int = 246534
+ n_positions: int = 256
+ n_embd: int = 1280
+ dff: int = 8192
+ n_layer: int = 48
+ n_head: int = 16
+ resid_pdrop: float | int = 0.1
+ embd_pdrop: float | int = 0.1
+ layer_norm_epsilon: float = 1e-6
+ initializer_range: float = 0.02
+ use_cache: bool = True
+ pad_token_id: int | None = None
+ bos_token_id: int | None = None
+ eos_token_id: int | list[int] | None = None
+ tie_word_embeddings: bool = True
+
+
+__all__ = ["CTRLConfig"]
diff --git a/third_party/transformers/src/transformers/models/ctrl/modeling_ctrl.py b/third_party/transformers/src/transformers/models/ctrl/modeling_ctrl.py
new file mode 100644
index 0000000000000000000000000000000000000000..1c2e6892e9a0d9781fc4e6d0be15cb9275abf493
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/ctrl/modeling_ctrl.py
@@ -0,0 +1,683 @@
+# Copyright 2018 Salesforce and HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# 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 CTRL model."""
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ... import initialization as init
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutput
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
+ auto_docstring,
+ logging,
+)
+from .configuration_ctrl import CTRLConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+def angle_defn(pos, i, d_model_size):
+ angle_rates = 1 / torch.pow(10000, (2 * (i // 2)) / d_model_size)
+ return pos * angle_rates
+
+
+def positional_encoding(position, d_model_size, dtype):
+ # create the sinusoidal pattern for the positional encoding
+ angle_rads = angle_defn(
+ torch.arange(position, dtype=torch.int64).to(dtype).unsqueeze(1),
+ torch.arange(d_model_size, dtype=torch.int64).to(dtype).unsqueeze(0),
+ d_model_size,
+ )
+
+ sines = torch.sin(angle_rads[:, 0::2])
+ cosines = torch.cos(angle_rads[:, 1::2])
+
+ pos_encoding = torch.cat([sines, cosines], dim=-1)
+ return pos_encoding
+
+
+def scaled_dot_product_attention(q, k, v, mask, attention_mask=None):
+ # calculate attention
+ matmul_qk = torch.matmul(q, k.permute(0, 1, 3, 2))
+
+ dk = k.shape[-1]
+ scaled_attention_logits = matmul_qk / np.sqrt(dk)
+
+ if mask is not None:
+ nd, ns = scaled_attention_logits.size(-2), scaled_attention_logits.size(-1)
+ scaled_attention_logits += mask[ns - nd : ns, :ns] * -1e4
+
+ if attention_mask is not None:
+ # Apply the attention mask
+ scaled_attention_logits = scaled_attention_logits + attention_mask
+
+ attention_weights = torch.softmax(scaled_attention_logits, dim=-1)
+
+ output = torch.matmul(attention_weights, v)
+
+ return output, attention_weights
+
+
+class MultiHeadAttention(nn.Module):
+ def __init__(self, d_model_size, num_heads, layer_idx=None):
+ super().__init__()
+ self.num_heads = num_heads
+ self.d_model_size = d_model_size
+ self.layer_idx = layer_idx
+
+ self.depth = int(d_model_size / self.num_heads)
+
+ self.Wq = nn.Linear(d_model_size, d_model_size)
+ self.Wk = nn.Linear(d_model_size, d_model_size)
+ self.Wv = nn.Linear(d_model_size, d_model_size)
+
+ self.dense = nn.Linear(d_model_size, d_model_size)
+
+ def split_into_heads(self, x, batch_size):
+ x = x.reshape(batch_size, -1, self.num_heads, self.depth)
+ return x.permute([0, 2, 1, 3])
+
+ def forward(
+ self,
+ v,
+ k,
+ q,
+ mask,
+ layer_past=None,
+ attention_mask=None,
+ use_cache=False,
+ output_attentions=False,
+ **kwargs,
+ ):
+ batch_size = q.shape[0]
+
+ q = self.Wq(q)
+ k = self.Wk(k)
+ v = self.Wv(v)
+
+ q = self.split_into_heads(q, batch_size)
+ k = self.split_into_heads(k, batch_size)
+ v = self.split_into_heads(v, batch_size)
+
+ if layer_past is not None:
+ k, v = layer_past.update(k, v, self.layer_idx)
+
+ output = scaled_dot_product_attention(q, k, v, mask, attention_mask)
+ scaled_attention = output[0].permute([0, 2, 1, 3])
+ attn = output[1]
+ original_size_attention = scaled_attention.reshape(batch_size, -1, self.d_model_size)
+ output = self.dense(original_size_attention)
+ return output, attn
+
+
+def point_wise_feed_forward_network(d_model_size, dff):
+ return nn.Sequential(nn.Linear(d_model_size, dff), nn.ReLU(), nn.Linear(dff, d_model_size))
+
+
+class EncoderLayer(nn.Module):
+ def __init__(self, d_model_size, num_heads, dff, rate=0.1, layer_idx=None):
+ super().__init__()
+
+ self.multi_head_attention = MultiHeadAttention(d_model_size, num_heads, layer_idx=layer_idx)
+ self.ffn = point_wise_feed_forward_network(d_model_size, dff)
+
+ self.layernorm1 = nn.LayerNorm(d_model_size, eps=1e-6)
+ self.layernorm2 = nn.LayerNorm(d_model_size, eps=1e-6)
+
+ self.dropout1 = nn.Dropout(rate)
+ self.dropout2 = nn.Dropout(rate)
+
+ def forward(
+ self,
+ x,
+ mask,
+ layer_past=None,
+ attention_mask=None,
+ use_cache=False,
+ output_attentions=False,
+ **kwargs,
+ ):
+ normed = self.layernorm1(x)
+ attn_outputs = self.multi_head_attention(
+ normed,
+ normed,
+ normed,
+ mask,
+ layer_past=layer_past,
+ attention_mask=attention_mask,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ )
+ attn_output = attn_outputs[0]
+ attn_output = self.dropout1(attn_output)
+ out1 = x + attn_output
+
+ out2 = self.layernorm2(out1)
+ ffn_output = self.ffn(out2)
+ ffn_output = self.dropout2(ffn_output)
+ out2 = out1 + ffn_output
+
+ outputs = (out2,) + attn_outputs[1:]
+ return outputs
+
+
+@auto_docstring
+class CTRLPreTrainedModel(PreTrainedModel):
+ config: CTRLConfig
+ base_model_prefix = "transformer"
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, CTRLModel):
+ init.copy_(
+ module.pos_encoding, positional_encoding(module.config.n_positions, module.d_model_size, torch.float)
+ )
+
+
+@auto_docstring
+class CTRLModel(CTRLPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.d_model_size = config.n_embd
+ self.num_layers = config.n_layer
+
+ self.w = nn.Embedding(config.vocab_size, config.n_embd)
+
+ self.dropout = nn.Dropout(config.embd_pdrop)
+ self.h = nn.ModuleList(
+ [
+ EncoderLayer(config.n_embd, config.n_head, config.dff, config.resid_pdrop, layer_idx=i)
+ for i in range(config.n_layer)
+ ]
+ )
+ self.layernorm = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+
+ self.register_buffer(
+ "pos_encoding", positional_encoding(config.n_positions, self.d_model_size, torch.float), persistent=False
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.w
+
+ def set_input_embeddings(self, new_embeddings):
+ self.w = new_embeddings
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs, # NOOP kwargs, for now
+ ) -> tuple[torch.Tensor] | BaseModelOutputWithPast:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, CTRLModel
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("Salesforce/ctrl")
+ >>> model = CTRLModel.from_pretrained("Salesforce/ctrl")
+
+ >>> # CTRL was trained with control codes as the first token
+ >>> inputs = tokenizer("Opinion My dog is cute", return_tensors="pt")
+ >>> assert inputs["input_ids"][0, 0].item() in tokenizer.control_codes.values()
+
+ >>> outputs = model(**inputs)
+
+ >>> last_hidden_states = outputs.last_hidden_state
+ >>> list(last_hidden_states.shape)
+ [1, 5, 1280]
+ ```"""
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+ 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.return_dict
+
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+ elif input_ids is not None:
+ self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
+ input_shape = input_ids.size()
+ input_ids = input_ids.view(-1, input_shape[-1])
+ batch_size = input_ids.shape[0]
+ elif inputs_embeds is not None:
+ input_shape = inputs_embeds.size()[:-1]
+ batch_size = inputs_embeds.shape[0]
+ else:
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ past_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+ if position_ids is None:
+ position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
+ position_ids = position_ids.unsqueeze(0)
+
+ # Attention mask.
+ if attention_mask is not None:
+ if batch_size <= 0:
+ raise ValueError("batch_size has to be defined and > 0")
+ attention_mask = attention_mask.view(batch_size, -1)
+ # We create a 3D attention mask from a 2D tensor mask.
+ # Sizes are [batch_size, 1, 1, to_seq_length]
+ # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
+ # this attention mask is more simple than the triangular masking of causal attention
+ # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
+ attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+
+ # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+ # masked positions, this operation will create a tensor which is 0.0 for
+ # positions we want to attend and the dtype's smallest value for masked positions.
+ # Since we are adding it to the raw scores before the softmax, this is
+ # effectively the same as removing these entirely.
+ attention_mask = attention_mask.to(dtype=self.dtype) # fp16 compatibility
+ attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min
+
+ if token_type_ids is not None:
+ token_type_ids = token_type_ids.view(-1, input_shape[-1])
+ token_type_embeds = self.w(token_type_ids)
+ token_type_embeds *= np.sqrt(self.d_model_size)
+ else:
+ token_type_embeds = 0
+
+ if inputs_embeds is None:
+ inputs_embeds = self.w(input_ids)
+ # inputs_embeds = embedded.unsqueeze(0) if len(input_ids.shape)<2 else embedded
+ seq_len = input_shape[-1]
+ mask = torch.triu(torch.ones(seq_len + past_length, seq_len + past_length), 1).to(device)
+
+ inputs_embeds *= np.sqrt(self.d_model_size)
+
+ # `self.pos_encoding` won't be sent to the correct device along the model, so we do it manually.
+ self.pos_encoding = self.pos_encoding.to(device)
+ pos_embeds = self.pos_encoding[position_ids, :]
+
+ hidden_states = inputs_embeds + pos_embeds + token_type_embeds
+
+ hidden_states = self.dropout(hidden_states)
+
+ all_hidden_states = () if output_hidden_states else None
+ all_attentions = () if output_attentions else None
+ for i, h in enumerate(self.h):
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+ outputs = h(
+ hidden_states,
+ mask,
+ layer_past=past_key_values,
+ attention_mask=attention_mask,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ )
+ hidden_states = outputs[0]
+ if output_attentions:
+ all_attentions += (outputs[1],)
+
+ hidden_states = self.layernorm(hidden_states)
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ if not return_dict:
+ return tuple(
+ v for v in [hidden_states, past_key_values, all_hidden_states, all_attentions] if v is not None
+ )
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ hidden_states=all_hidden_states,
+ attentions=all_attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The CTRL Model transformer with a language modeling head on top (linear layer with weights tied to the input
+ embeddings).
+ """
+)
+class CTRLLMHeadModel(CTRLPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "transformer.w.weight"}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.transformer = CTRLModel(config)
+ self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=True)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs,
+ ) -> tuple[torch.Tensor] | CausalLMOutputWithPast:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+ `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+ are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+
+ Example:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoTokenizer, CTRLLMHeadModel
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("Salesforce/ctrl")
+ >>> model = CTRLLMHeadModel.from_pretrained("Salesforce/ctrl")
+
+ >>> # CTRL was trained with control codes as the first token
+ >>> inputs = tokenizer("Wikipedia The llama is", return_tensors="pt")
+ >>> assert inputs["input_ids"][0, 0].item() in tokenizer.control_codes.values()
+
+ >>> sequence_ids = model.generate(inputs["input_ids"])
+ >>> sequences = tokenizer.batch_decode(sequence_ids)
+ >>> sequences
+ ['Wikipedia The llama is a member of the family Bovidae. It is native to the Andes of Peru,']
+
+ >>> outputs = model(**inputs, labels=inputs["input_ids"])
+ >>> round(outputs.loss.item(), 2)
+ 9.21
+
+ >>> list(outputs.logits.shape)
+ [1, 5, 246534]
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ transformer_outputs = self.transformer(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = transformer_outputs[0]
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits,
+ labels,
+ vocab_size=self.config.vocab_size,
+ **kwargs,
+ )
+
+ if not return_dict:
+ output = (logits,) + transformer_outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=transformer_outputs.past_key_values,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+ def prepare_inputs_for_generation(
+ self, input_ids, past_key_values=None, use_cache=None, is_first_iteration=False, **kwargs
+ ):
+ # Overwritten -- `token_type_ids` are created in custom way inside model`
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ # token_type_ids are computed on CTRLModel.forward()
+ model_inputs.pop("token_type_ids", None)
+
+ return model_inputs
+
+
+@auto_docstring(
+ custom_intro="""
+ The CTRL Model transformer with a sequence classification head on top (linear layer).
+ [`CTRLForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+ (e.g. GPT-2) do. Since it does classification on the last token, it requires to know the position of the last
+ token. If a `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in
+ each row. If no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot
+ guess the padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last
+ value in each row of the batch).
+ """
+)
+class CTRLForSequenceClassification(CTRLPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.transformer = CTRLModel(config)
+ self.classifier = nn.Linear(config.n_embd, self.num_labels, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor] | SequenceClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+
+ Example of single-label classification:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoTokenizer, CTRLForSequenceClassification
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("Salesforce/ctrl")
+ >>> model = CTRLForSequenceClassification.from_pretrained("Salesforce/ctrl")
+
+ >>> # CTRL was trained with control codes as the first token
+ >>> inputs = tokenizer("Opinion My dog is cute", return_tensors="pt")
+ >>> assert inputs["input_ids"][0, 0].item() in tokenizer.control_codes.values()
+
+ >>> with torch.no_grad():
+ ... logits = model(**inputs).logits
+
+ >>> predicted_class_id = logits.argmax().item()
+ >>> model.config.id2label[predicted_class_id]
+ 'LABEL_0'
+ ```
+
+ ```python
+ >>> import torch
+
+ >>> torch.manual_seed(42) # doctest: +IGNORE_RESULT
+ >>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`
+ >>> num_labels = len(model.config.id2label)
+ >>> model = CTRLForSequenceClassification.from_pretrained("Salesforce/ctrl", num_labels=num_labels)
+
+ >>> labels = torch.tensor(1)
+ >>> loss = model(**inputs, labels=labels).loss
+ >>> round(loss.item(), 2)
+ 0.93
+ ```
+
+ Example of multi-label classification:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoTokenizer, CTRLForSequenceClassification
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("Salesforce/ctrl")
+ >>> model = CTRLForSequenceClassification.from_pretrained(
+ ... "Salesforce/ctrl", problem_type="multi_label_classification"
+ ... )
+
+ >>> # CTRL was trained with control codes as the first token
+ >>> inputs = tokenizer("Opinion My dog is cute", return_tensors="pt")
+ >>> assert inputs["input_ids"][0, 0].item() in tokenizer.control_codes.values()
+
+ >>> with torch.no_grad():
+ ... logits = model(**inputs).logits
+
+ >>> predicted_class_id = logits.argmax().item()
+ >>> model.config.id2label[predicted_class_id]
+ 'LABEL_0'
+ ```
+
+ ```python
+ >>> # To train a model on `num_labels` classes, you can pass `num_labels=num_labels` to `.from_pretrained(...)`
+ >>> num_labels = len(model.config.id2label)
+ >>> model = CTRLForSequenceClassification.from_pretrained("Salesforce/ctrl", num_labels=num_labels)
+
+ >>> num_labels = len(model.config.id2label)
+ >>> labels = torch.nn.functional.one_hot(torch.tensor([predicted_class_id]), num_classes=num_labels).to(
+ ... torch.float
+ ... )
+ >>> loss = model(**inputs, labels=labels).loss
+ >>> loss.backward() # doctest: +IGNORE_RESULT
+ ```"""
+
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ transformer_outputs = self.transformer(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = transformer_outputs[0]
+ logits = self.classifier(hidden_states)
+
+ if input_ids is not None:
+ batch_size, sequence_length = input_ids.shape[:2]
+ else:
+ batch_size, sequence_length = inputs_embeds.shape[:2]
+
+ if self.config.pad_token_id is None and batch_size != 1:
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+ if self.config.pad_token_id is None:
+ last_non_pad_token = -1
+ elif input_ids is not None:
+ # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
+ non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+ token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+ last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+ else:
+ last_non_pad_token = -1
+ logger.warning_once(
+ f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+ "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+ )
+
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]
+
+ loss = None
+ if labels is not None:
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(pooled_logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(pooled_logits, labels)
+ if not return_dict:
+ output = (pooled_logits,) + transformer_outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return SequenceClassifierOutput(
+ loss=loss,
+ logits=pooled_logits,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+
+__all__ = ["CTRLForSequenceClassification", "CTRLLMHeadModel", "CTRLModel", "CTRLPreTrainedModel"]
diff --git a/third_party/transformers/src/transformers/models/ctrl/tokenization_ctrl.py b/third_party/transformers/src/transformers/models/ctrl/tokenization_ctrl.py
new file mode 100644
index 0000000000000000000000000000000000000000..6ca2b83fba22cce4a5a140a896aad8ed5069b4f4
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/ctrl/tokenization_ctrl.py
@@ -0,0 +1,226 @@
+# Copyright 2018 Salesforce and The HuggingFace Inc. team.
+#
+# 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.
+"""Tokenization classes for Salesforce CTRL."""
+
+import json
+
+import regex as re
+
+from ...tokenization_python import PreTrainedTokenizer
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {
+ "vocab_file": "vocab.json",
+ "merges_file": "merges.txt",
+}
+
+
+CONTROL_CODES = {
+ "Pregnancy": 168629,
+ "Christianity": 7675,
+ "Explain": 106423,
+ "Fitness": 63440,
+ "Saving": 63163,
+ "Ask": 27171,
+ "Ass": 95985,
+ "Joke": 163509,
+ "Questions": 45622,
+ "Thoughts": 49605,
+ "Retail": 52342,
+ "Feminism": 164338,
+ "Writing": 11992,
+ "Atheism": 192263,
+ "Netflix": 48616,
+ "Computing": 39639,
+ "Opinion": 43213,
+ "Alone": 44967,
+ "Funny": 58917,
+ "Gaming": 40358,
+ "Human": 4088,
+ "India": 1331,
+ "Joker": 77138,
+ "Diet": 36206,
+ "Legal": 11859,
+ "Norman": 4939,
+ "Tip": 72689,
+ "Weight": 52343,
+ "Movies": 46273,
+ "Running": 23425,
+ "Science": 2090,
+ "Horror": 37793,
+ "Confession": 60572,
+ "Finance": 12250,
+ "Politics": 16360,
+ "Scary": 191985,
+ "Support": 12654,
+ "Technologies": 32516,
+ "Teenage": 66160,
+ "Event": 32769,
+ "Learned": 67460,
+ "Notion": 182770,
+ "Wikipedia": 37583,
+ "Books": 6665,
+ "Extract": 76050,
+ "Confessions": 102701,
+ "Conspiracy": 75932,
+ "Links": 63674,
+ "Narcissus": 150425,
+ "Relationship": 54766,
+ "Relationships": 134796,
+ "Reviews": 41671,
+ "News": 4256,
+ "Translation": 26820,
+ "multilingual": 128406,
+}
+
+
+def get_pairs(word):
+ """
+ Return set of symbol pairs in a word.
+
+ Word is represented as tuple of symbols (symbols being variable-length strings).
+ """
+ pairs = set()
+ prev_char = word[0]
+ for char in word[1:]:
+ pairs.add((prev_char, char))
+ prev_char = char
+
+ pairs = set(pairs)
+ return pairs
+
+
+class CTRLTokenizer(PreTrainedTokenizer):
+ """
+ Construct a CTRL tokenizer. Based on Byte-Pair-Encoding.
+
+ This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
+ this superclass for more information regarding those methods.
+
+ Args:
+ vocab_file (`str`):
+ Path to the vocabulary file.
+ merges_file (`str`):
+ Path to the merges file.
+ unk_token (`str`, *optional*, defaults to `""`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ control_codes = CONTROL_CODES
+
+ def __init__(self, vocab_file, merges_file, unk_token="", **kwargs):
+ with open(vocab_file, encoding="utf-8") as vocab_handle:
+ self.encoder = json.load(vocab_handle)
+ self.decoder = {v: k for k, v in self.encoder.items()}
+ with open(merges_file, encoding="utf-8") as merges_handle:
+ merges = merges_handle.read().split("\n")[1:-1]
+ merges = [tuple(merge.split()) for merge in merges]
+ self.bpe_ranks = dict(zip(merges, range(len(merges))))
+ self.cache = {}
+ self.add_bpe_version_header = True
+ super().__init__(
+ unk_token=unk_token,
+ token_type_ids_pattern="all_zeros",
+ token_type_ids_include_special_tokens=True,
+ special_tokens_pattern="none",
+ **kwargs,
+ )
+
+ @property
+ def vocab_size(self):
+ return len(self.encoder)
+
+ def get_vocab(self):
+ return dict(self.encoder, **self.added_tokens_encoder)
+
+ def bpe(self, token):
+ if token in self.cache:
+ return self.cache[token]
+ word = tuple(token)
+ word = tuple(list(word[:-1]) + [word[-1] + ""])
+ pairs = get_pairs(word)
+
+ if not pairs:
+ return token
+
+ while True:
+ bigram = min(pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf")))
+ if bigram not in self.bpe_ranks:
+ break
+ first, second = bigram
+ new_word = []
+ i = 0
+ while i < len(word):
+ try:
+ j = word.index(first, i)
+ except ValueError:
+ new_word.extend(word[i:])
+ break
+ else:
+ new_word.extend(word[i:j])
+ i = j
+
+ if word[i] == first and i < len(word) - 1 and word[i + 1] == second:
+ new_word.append(first + second)
+ i += 2
+ else:
+ new_word.append(word[i])
+ i += 1
+ new_word = tuple(new_word)
+ word = new_word
+ if len(word) == 1:
+ break
+ else:
+ pairs = get_pairs(word)
+ word = "@@ ".join(word)
+ word = word[:-4]
+ self.cache[token] = word
+ return word
+
+ def _tokenize(self, text):
+ """Tokenize a string."""
+ split_tokens = []
+
+ words = re.findall(r"\S+\n?", text)
+
+ for token in words:
+ split_tokens.extend(list(self.bpe(token).split(" ")))
+ return split_tokens
+
+ def _convert_token_to_id(self, token):
+ """Converts a token (str) in an id using the vocab."""
+ return self.encoder.get(token, self.encoder.get(self.unk_token))
+
+ def _convert_id_to_token(self, index):
+ """Converts an index (integer) in a token (str) using the vocab."""
+ return self.decoder.get(index, self.unk_token)
+
+ def convert_tokens_to_string(self, tokens):
+ """Converts a sequence of tokens (string) in a single string."""
+ out_string = " ".join(tokens).replace("@@ ", "").strip()
+ return out_string
+
+ # def decode(self, token_ids, skip_special_tokens=False, clean_up_tokenization_spaces=True):
+ # filtered_tokens = ' '.join(self.convert_ids_to_tokens(token_ids, skip_special_tokens=skip_special_tokens))
+ # tokens_generated_so_far = re.sub('(@@ )', '', string=filtered_tokens)
+ # tokens_generated_so_far = re.sub('(@@ ?$)', '', string=tokens_generated_so_far)
+ # return ''.join(tokens_generated_so_far)
+
+
+__all__ = ["CTRLTokenizer"]
diff --git a/third_party/transformers/src/transformers/models/dinat/__init__.py b/third_party/transformers/src/transformers/models/dinat/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..b64cdbb3c7eb0467f6112225b8c0d9e1f65f9e99
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/dinat/__init__.py
@@ -0,0 +1,27 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_dinat import *
+ from .modeling_dinat import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/dinat/configuration_dinat.py b/third_party/transformers/src/transformers/models/dinat/configuration_dinat.py
new file mode 100644
index 0000000000000000000000000000000000000000..a5c48765e4fcf2fda6768b8120ba093143dee07d
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/dinat/configuration_dinat.py
@@ -0,0 +1,85 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Dilated Neighborhood Attention Transformer model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...backbone_utils import BackboneConfigMixin
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="shi-labs/dinat-mini-in1k-224")
+@strict
+class DinatConfig(BackboneConfigMixin, PreTrainedConfig):
+ r"""
+ dilations (`list[list[int]]`, *optional*, defaults to `[[1, 8, 1], [1, 4, 1, 4], [1, 2, 1, 2, 1, 2], [1, 1, 1, 1, 1]]`):
+ Dilation value of each NA layer in the Transformer encoder.
+
+ Example:
+
+ ```python
+ >>> from transformers import DinatConfig, DinatModel
+
+ >>> # Initializing a Dinat shi-labs/dinat-mini-in1k-224 style configuration
+ >>> configuration = DinatConfig()
+
+ >>> # Initializing a model (with random weights) from the shi-labs/dinat-mini-in1k-224 style configuration
+ >>> model = DinatModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "dinat"
+
+ attribute_map = {
+ "num_attention_heads": "num_heads",
+ "num_hidden_layers": "num_layers",
+ }
+
+ patch_size: int | list[int] | tuple[int, int] = 4
+ num_channels: int = 3
+ embed_dim: int = 64
+ depths: list[int] | tuple[int, ...] = (3, 4, 6, 5)
+ num_heads: list[int] | tuple[int, ...] = (2, 4, 8, 16)
+ kernel_size: int = 7
+ dilations: list | tuple | None = None
+ mlp_ratio: float = 3.0
+ qkv_bias: bool = True
+ hidden_dropout_prob: float | int = 0.0
+ attention_probs_dropout_prob: float | int = 0.0
+ drop_path_rate: float | int = 0.1
+ hidden_act: str = "gelu"
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-5
+ layer_scale_init_value: float = 0.0
+ _out_features: list[str] | None = None
+ _out_indices: list[int] | None = None
+
+ def __post_init__(self, **kwargs):
+ self.num_layers = len(self.depths)
+ self.dilations = self.dilations or [[1, 8, 1], [1, 4, 1, 4], [1, 2, 1, 2, 1, 2], [1, 1, 1, 1, 1]]
+
+ # we set the hidden_size attribute in order to make Dinat work with VisionEncoderDecoderModel
+ # this indicates the channel dimension after the last stage of the model
+ self.hidden_size = int(self.embed_dim * 2 ** (len(self.depths) - 1))
+ self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, len(self.depths) + 1)]
+ self.set_output_features_output_indices(
+ out_indices=kwargs.pop("out_indices", None), out_features=kwargs.pop("out_features", None)
+ )
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["DinatConfig"]
diff --git a/third_party/transformers/src/transformers/models/dinat/modeling_dinat.py b/third_party/transformers/src/transformers/models/dinat/modeling_dinat.py
new file mode 100644
index 0000000000000000000000000000000000000000..783231ff0192419b280bd83133d4f57bdc27ab12
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/dinat/modeling_dinat.py
@@ -0,0 +1,816 @@
+# Copyright 2022 SHI Labs and The HuggingFace Inc. team. All rights reserved.
+#
+# 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 Dilated Neighborhood Attention Transformer model."""
+
+import math
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...backbone_utils import BackboneMixin, filter_output_hidden_states
+from ...modeling_outputs import BackboneOutput
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
+ ModelOutput,
+ OptionalDependencyNotAvailable,
+ auto_docstring,
+ is_natten_available,
+ logging,
+ requires_backends,
+)
+from ...utils.generic import can_return_tuple
+from .configuration_dinat import DinatConfig
+
+
+if is_natten_available():
+ from natten.functional import natten2dav, natten2dqkrpb
+else:
+
+ def natten2dqkrpb(*args, **kwargs):
+ raise OptionalDependencyNotAvailable()
+
+ def natten2dav(*args, **kwargs):
+ raise OptionalDependencyNotAvailable()
+
+
+logger = logging.get_logger(__name__)
+
+
+# drop_path and DinatDropPath are from the timm library.
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Dinat encoder's outputs, with potential hidden states and attentions.
+ """
+)
+class DinatEncoderOutput(ModelOutput):
+ r"""
+ reshaped_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 + one for the output of each stage) of
+ shape `(batch_size, hidden_size, height, width)`.
+
+ Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+ include the spatial dimensions.
+ """
+
+ last_hidden_state: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+ reshaped_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Dinat model's outputs that also contains a pooling of the last hidden states.
+ """
+)
+class DinatModelOutput(ModelOutput):
+ r"""
+ pooler_output (`torch.FloatTensor` of shape `(batch_size, hidden_size)`, *optional*, returned when `add_pooling_layer=True` is passed):
+ Average pooling of the last layer hidden-state.
+ reshaped_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 + one for the output of each stage) of
+ shape `(batch_size, hidden_size, height, width)`.
+
+ Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+ include the spatial dimensions.
+ """
+
+ last_hidden_state: torch.FloatTensor | None = None
+ pooler_output: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+ reshaped_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Dinat outputs for image classification.
+ """
+)
+class DinatImageClassifierOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided):
+ Classification (or regression if config.num_labels==1) loss.
+ logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels)`):
+ Classification (or regression if config.num_labels==1) scores (before SoftMax).
+ reshaped_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 + one for the output of each stage) of
+ shape `(batch_size, hidden_size, height, width)`.
+
+ Hidden-states of the model at the output of each layer plus the initial embedding outputs reshaped to
+ include the spatial dimensions.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+ reshaped_hidden_states: tuple[torch.FloatTensor, ...] | None = None
+
+
+class DinatEmbeddings(nn.Module):
+ """
+ Construct the patch and position embeddings.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+
+ self.patch_embeddings = DinatPatchEmbeddings(config)
+
+ self.norm = nn.LayerNorm(config.embed_dim)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, pixel_values: torch.FloatTensor | None) -> tuple[torch.Tensor]:
+ embeddings = self.patch_embeddings(pixel_values)
+ embeddings = self.norm(embeddings)
+
+ embeddings = self.dropout(embeddings)
+
+ return embeddings
+
+
+class DinatPatchEmbeddings(nn.Module):
+ """
+ This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
+ `hidden_states` (patch embeddings) of shape `(batch_size, height, width, hidden_size)` to be consumed by a
+ Transformer.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ patch_size = config.patch_size
+ num_channels, hidden_size = config.num_channels, config.embed_dim
+ self.num_channels = num_channels
+
+ if patch_size == 4:
+ pass
+ else:
+ # TODO: Support arbitrary patch sizes.
+ raise ValueError("Dinat only supports patch size of 4 at the moment.")
+
+ self.projection = nn.Sequential(
+ nn.Conv2d(self.num_channels, hidden_size // 2, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)),
+ nn.Conv2d(hidden_size // 2, hidden_size, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)),
+ )
+
+ def forward(self, pixel_values: torch.FloatTensor | None) -> torch.Tensor:
+ _, num_channels, height, width = pixel_values.shape
+ if num_channels != self.num_channels:
+ raise ValueError(
+ "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+ )
+ embeddings = self.projection(pixel_values)
+ embeddings = embeddings.permute(0, 2, 3, 1)
+
+ return embeddings
+
+
+class DinatDownsampler(nn.Module):
+ """
+ Convolutional Downsampling Layer.
+
+ Args:
+ dim (`int`):
+ Number of input channels.
+ norm_layer (`nn.Module`, *optional*, defaults to `nn.LayerNorm`):
+ Normalization layer class.
+ """
+
+ def __init__(self, dim: int, norm_layer: nn.Module = nn.LayerNorm) -> None:
+ super().__init__()
+ self.dim = dim
+ self.reduction = nn.Conv2d(dim, 2 * dim, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)
+ self.norm = norm_layer(2 * dim)
+
+ def forward(self, input_feature: torch.Tensor) -> torch.Tensor:
+ input_feature = self.reduction(input_feature.permute(0, 3, 1, 2)).permute(0, 2, 3, 1)
+ input_feature = self.norm(input_feature)
+ return input_feature
+
+
+# Copied from transformers.models.beit.modeling_beit.drop_path
+def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
+ """
+ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+ """
+ if drop_prob == 0.0 or not training:
+ return input
+ keep_prob = 1 - drop_prob
+ shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
+ random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
+ random_tensor.floor_() # binarize
+ output = input.div(keep_prob) * random_tensor
+ return output
+
+
+# Copied from transformers.models.beit.modeling_beit.BeitDropPath with Beit->Dinat
+class DinatDropPath(nn.Module):
+ """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+ def __init__(self, drop_prob: float | None = None) -> None:
+ super().__init__()
+ self.drop_prob = drop_prob
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ return drop_path(hidden_states, self.drop_prob, self.training)
+
+ def extra_repr(self) -> str:
+ return f"p={self.drop_prob}"
+
+
+class NeighborhoodAttention(nn.Module):
+ def __init__(self, config, dim, num_heads, kernel_size, dilation):
+ super().__init__()
+ if dim % num_heads != 0:
+ raise ValueError(
+ f"The hidden size ({dim}) is not a multiple of the number of attention heads ({num_heads})"
+ )
+
+ self.num_attention_heads = num_heads
+ self.attention_head_size = int(dim / num_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+ self.kernel_size = kernel_size
+ self.dilation = dilation
+
+ # rpb is learnable relative positional biases; same concept is used Swin.
+ self.rpb = nn.Parameter(torch.zeros(num_heads, (2 * self.kernel_size - 1), (2 * self.kernel_size - 1)))
+
+ self.query = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
+ self.key = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
+ self.value = nn.Linear(self.all_head_size, self.all_head_size, bias=config.qkv_bias)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ output_attentions: bool | None = False,
+ ) -> tuple[torch.Tensor]:
+ batch_size, seq_length, _ = hidden_states.shape
+ query_layer = (
+ self.query(hidden_states)
+ .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)
+ .transpose(1, 2)
+ )
+ key_layer = (
+ self.key(hidden_states)
+ .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)
+ .transpose(1, 2)
+ )
+ value_layer = (
+ self.value(hidden_states)
+ .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)
+ .transpose(1, 2)
+ )
+
+ # Apply the scale factor before computing attention weights. It's usually more efficient because
+ # attention weights are typically a bigger tensor compared to query.
+ # It gives identical results because scalars are commutable in matrix multiplication.
+ query_layer = query_layer / math.sqrt(self.attention_head_size)
+
+ # Compute NA between "query" and "key" to get the raw attention scores, and add relative positional biases.
+ attention_scores = natten2dqkrpb(query_layer, key_layer, self.rpb, self.kernel_size, self.dilation)
+
+ # Normalize the attention scores to probabilities.
+ attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+
+ # This is actually dropping out entire tokens to attend to, which might
+ # seem a bit unusual, but is taken from the original Transformer paper.
+ attention_probs = self.dropout(attention_probs)
+
+ context_layer = natten2dav(attention_probs, value_layer, self.kernel_size, self.dilation)
+ context_layer = context_layer.permute(0, 2, 3, 1, 4).contiguous()
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+ context_layer = context_layer.view(new_context_layer_shape)
+
+ outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+
+ return outputs
+
+
+class NeighborhoodAttentionOutput(nn.Module):
+ def __init__(self, config, dim):
+ super().__init__()
+ self.dense = nn.Linear(dim, dim)
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+
+ return hidden_states
+
+
+class NeighborhoodAttentionModule(nn.Module):
+ def __init__(self, config, dim, num_heads, kernel_size, dilation):
+ super().__init__()
+ self.self = NeighborhoodAttention(config, dim, num_heads, kernel_size, dilation)
+ self.output = NeighborhoodAttentionOutput(config, dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ output_attentions: bool | None = False,
+ ) -> tuple[torch.Tensor]:
+ self_outputs = self.self(hidden_states, output_attentions)
+ attention_output = self.output(self_outputs[0], hidden_states)
+ outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
+ return outputs
+
+
+class DinatIntermediate(nn.Module):
+ def __init__(self, config, dim):
+ super().__init__()
+ self.dense = nn.Linear(dim, int(config.mlp_ratio * dim))
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+class DinatOutput(nn.Module):
+ def __init__(self, config, dim):
+ super().__init__()
+ self.dense = nn.Linear(int(config.mlp_ratio * dim), dim)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ return hidden_states
+
+
+class DinatLayer(nn.Module):
+ def __init__(self, config, dim, num_heads, dilation, drop_path_rate=0.0):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.kernel_size = config.kernel_size
+ self.dilation = dilation
+ self.window_size = self.kernel_size * self.dilation
+ self.layernorm_before = nn.LayerNorm(dim, eps=config.layer_norm_eps)
+ self.attention = NeighborhoodAttentionModule(
+ config, dim, num_heads, kernel_size=self.kernel_size, dilation=self.dilation
+ )
+ self.drop_path = DinatDropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
+ self.layernorm_after = nn.LayerNorm(dim, eps=config.layer_norm_eps)
+ self.intermediate = DinatIntermediate(config, dim)
+ self.output = DinatOutput(config, dim)
+ self.layer_scale_parameters = (
+ nn.Parameter(config.layer_scale_init_value * torch.ones((2, dim)), requires_grad=True)
+ if config.layer_scale_init_value > 0
+ else None
+ )
+
+ def maybe_pad(self, hidden_states, height, width):
+ window_size = self.window_size
+ pad_values = (0, 0, 0, 0, 0, 0)
+ if height < window_size or width < window_size:
+ pad_l = pad_t = 0
+ pad_r = max(0, window_size - width)
+ pad_b = max(0, window_size - height)
+ pad_values = (0, 0, pad_l, pad_r, pad_t, pad_b)
+ hidden_states = nn.functional.pad(hidden_states, pad_values)
+ return hidden_states, pad_values
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ output_attentions: bool | None = False,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ batch_size, height, width, channels = hidden_states.size()
+ shortcut = hidden_states
+
+ hidden_states = self.layernorm_before(hidden_states)
+ # pad hidden_states if they are smaller than kernel size x dilation
+ hidden_states, pad_values = self.maybe_pad(hidden_states, height, width)
+
+ _, height_pad, width_pad, _ = hidden_states.shape
+
+ attention_outputs = self.attention(hidden_states, output_attentions=output_attentions)
+
+ attention_output = attention_outputs[0]
+
+ was_padded = pad_values[3] > 0 or pad_values[5] > 0
+ if was_padded:
+ attention_output = attention_output[:, :height, :width, :].contiguous()
+
+ if self.layer_scale_parameters is not None:
+ attention_output = self.layer_scale_parameters[0] * attention_output
+
+ hidden_states = shortcut + self.drop_path(attention_output)
+
+ layer_output = self.layernorm_after(hidden_states)
+ layer_output = self.output(self.intermediate(layer_output))
+
+ if self.layer_scale_parameters is not None:
+ layer_output = self.layer_scale_parameters[1] * layer_output
+
+ layer_output = hidden_states + self.drop_path(layer_output)
+
+ layer_outputs = (layer_output, attention_outputs[1]) if output_attentions else (layer_output,)
+ return layer_outputs
+
+
+class DinatStage(nn.Module):
+ def __init__(self, config, dim, depth, num_heads, dilations, drop_path_rate, downsample):
+ super().__init__()
+ self.config = config
+ self.dim = dim
+ self.layers = nn.ModuleList(
+ [
+ DinatLayer(
+ config=config,
+ dim=dim,
+ num_heads=num_heads,
+ dilation=dilations[i],
+ drop_path_rate=drop_path_rate[i],
+ )
+ for i in range(depth)
+ ]
+ )
+
+ # patch merging layer
+ if downsample is not None:
+ self.downsample = downsample(dim=dim, norm_layer=nn.LayerNorm)
+ else:
+ self.downsample = None
+
+ self.pointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ output_attentions: bool | None = False,
+ ) -> tuple[torch.Tensor]:
+ _, height, width, _ = hidden_states.size()
+ for i, layer_module in enumerate(self.layers):
+ layer_outputs = layer_module(hidden_states, output_attentions)
+ hidden_states = layer_outputs[0]
+
+ hidden_states_before_downsampling = hidden_states
+ if self.downsample is not None:
+ hidden_states = self.downsample(hidden_states_before_downsampling)
+
+ stage_outputs = (hidden_states, hidden_states_before_downsampling)
+
+ if output_attentions:
+ stage_outputs += layer_outputs[1:]
+ return stage_outputs
+
+
+class DinatEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.num_levels = len(config.depths)
+ self.config = config
+ dpr = [x.item() for x in torch.linspace(0, config.drop_path_rate, sum(config.depths), device="cpu")]
+ self.levels = nn.ModuleList(
+ [
+ DinatStage(
+ config=config,
+ dim=int(config.embed_dim * 2**i_layer),
+ depth=config.depths[i_layer],
+ num_heads=config.num_heads[i_layer],
+ dilations=config.dilations[i_layer],
+ drop_path_rate=dpr[sum(config.depths[:i_layer]) : sum(config.depths[: i_layer + 1])],
+ downsample=DinatDownsampler if (i_layer < self.num_levels - 1) else None,
+ )
+ for i_layer in range(self.num_levels)
+ ]
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ output_attentions: bool | None = False,
+ output_hidden_states: bool | None = False,
+ output_hidden_states_before_downsampling: bool | None = False,
+ return_dict: bool | None = True,
+ ) -> tuple | DinatEncoderOutput:
+ all_hidden_states = () if output_hidden_states else None
+ all_reshaped_hidden_states = () if output_hidden_states else None
+ all_self_attentions = () if output_attentions else None
+
+ if output_hidden_states:
+ # rearrange b h w c -> b c h w
+ reshaped_hidden_state = hidden_states.permute(0, 3, 1, 2)
+ all_hidden_states += (hidden_states,)
+ all_reshaped_hidden_states += (reshaped_hidden_state,)
+
+ for i, layer_module in enumerate(self.levels):
+ layer_outputs = layer_module(hidden_states, output_attentions)
+
+ hidden_states = layer_outputs[0]
+ hidden_states_before_downsampling = layer_outputs[1]
+
+ if output_hidden_states and output_hidden_states_before_downsampling:
+ # rearrange b h w c -> b c h w
+ reshaped_hidden_state = hidden_states_before_downsampling.permute(0, 3, 1, 2)
+ all_hidden_states += (hidden_states_before_downsampling,)
+ all_reshaped_hidden_states += (reshaped_hidden_state,)
+ elif output_hidden_states and not output_hidden_states_before_downsampling:
+ # rearrange b h w c -> b c h w
+ reshaped_hidden_state = hidden_states.permute(0, 3, 1, 2)
+ all_hidden_states += (hidden_states,)
+ all_reshaped_hidden_states += (reshaped_hidden_state,)
+
+ if output_attentions:
+ all_self_attentions += layer_outputs[2:]
+
+ if not return_dict:
+ return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
+
+ return DinatEncoderOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ reshaped_hidden_states=all_reshaped_hidden_states,
+ )
+
+
+@auto_docstring
+class DinatPreTrainedModel(PreTrainedModel):
+ config: DinatConfig
+ base_model_prefix = "dinat"
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+
+
+@auto_docstring
+class DinatModel(DinatPreTrainedModel):
+ def __init__(self, config, add_pooling_layer=True):
+ r"""
+ add_pooling_layer (bool, *optional*, defaults to `True`):
+ Whether to add a pooling layer
+ """
+ super().__init__(config)
+
+ requires_backends(self, ["natten"])
+
+ self.config = config
+ self.num_levels = len(config.depths)
+ self.num_features = int(config.embed_dim * 2 ** (self.num_levels - 1))
+
+ self.embeddings = DinatEmbeddings(config)
+ self.encoder = DinatEncoder(config)
+
+ self.layernorm = nn.LayerNorm(self.num_features, eps=config.layer_norm_eps)
+ self.pooler = nn.AdaptiveAvgPool1d(1) if add_pooling_layer else None
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embeddings.patch_embeddings
+
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | DinatModelOutput:
+ 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.return_dict
+
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ embedding_output = self.embeddings(pixel_values)
+
+ encoder_outputs = self.encoder(
+ embedding_output,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ sequence_output = encoder_outputs[0]
+ sequence_output = self.layernorm(sequence_output)
+
+ pooled_output = None
+ if self.pooler is not None:
+ pooled_output = self.pooler(sequence_output.flatten(1, 2).transpose(1, 2))
+ pooled_output = torch.flatten(pooled_output, 1)
+
+ if not return_dict:
+ output = (sequence_output, pooled_output) + encoder_outputs[1:]
+
+ return output
+
+ return DinatModelOutput(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ reshaped_hidden_states=encoder_outputs.reshaped_hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Dinat Model transformer with an image classification head on top (a linear layer on top of the final hidden state
+ of the [CLS] token) e.g. for ImageNet.
+ """
+)
+class DinatForImageClassification(DinatPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ requires_backends(self, ["natten"])
+
+ self.num_labels = config.num_labels
+ self.dinat = DinatModel(config)
+
+ # Classifier head
+ self.classifier = (
+ nn.Linear(self.dinat.num_features, config.num_labels) if config.num_labels > 0 else nn.Identity()
+ )
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | DinatImageClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the image classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.dinat(
+ pixel_values,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ pooled_output = outputs[1]
+
+ logits = self.classifier(pooled_output)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(labels, logits, self.config)
+
+ if not return_dict:
+ output = (logits,) + outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return DinatImageClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ reshaped_hidden_states=outputs.reshaped_hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ NAT backbone, to be used with frameworks like DETR and MaskFormer.
+ """
+)
+class DinatBackbone(BackboneMixin, DinatPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ requires_backends(self, ["natten"])
+
+ self.embeddings = DinatEmbeddings(config)
+ self.encoder = DinatEncoder(config)
+ self.num_features = [config.embed_dim] + [int(config.embed_dim * 2**i) for i in range(len(config.depths))]
+
+ # Add layer norms to hidden states of out_features
+ hidden_states_norms = {}
+ for stage, num_channels in zip(self.out_features, self.channels):
+ hidden_states_norms[stage] = nn.LayerNorm(num_channels)
+ self.hidden_states_norms = nn.ModuleDict(hidden_states_norms)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embeddings.patch_embeddings
+
+ @can_return_tuple
+ @filter_output_hidden_states
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.Tensor,
+ output_hidden_states: bool | None = None,
+ output_attentions: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> BackboneOutput:
+ r"""
+ Examples:
+
+ ```python
+ >>> from transformers import AutoImageProcessor, AutoBackbone
+ >>> import torch
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> processor = AutoImageProcessor.from_pretrained("shi-labs/nat-mini-in1k-224")
+ >>> model = AutoBackbone.from_pretrained(
+ ... "shi-labs/nat-mini-in1k-224", out_features=["stage1", "stage2", "stage3", "stage4"]
+ ... )
+
+ >>> inputs = processor(image, return_tensors="pt")
+
+ >>> outputs = model(**inputs)
+
+ >>> feature_maps = outputs.feature_maps
+ >>> list(feature_maps[-1].shape)
+ [1, 512, 7, 7]
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
+ embedding_output = self.embeddings(pixel_values)
+
+ outputs = self.encoder(
+ embedding_output,
+ output_attentions=output_attentions,
+ output_hidden_states=True,
+ output_hidden_states_before_downsampling=True,
+ return_dict=True,
+ )
+
+ hidden_states = outputs.reshaped_hidden_states
+
+ feature_maps = ()
+ for stage, hidden_state in zip(self.stage_names, hidden_states):
+ if stage in self.out_features:
+ batch_size, num_channels, height, width = hidden_state.shape
+ hidden_state = hidden_state.permute(0, 2, 3, 1).contiguous()
+ hidden_state = hidden_state.view(batch_size, height * width, num_channels)
+ hidden_state = self.hidden_states_norms[stage](hidden_state)
+ hidden_state = hidden_state.view(batch_size, height, width, num_channels)
+ hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous()
+ feature_maps += (hidden_state,)
+
+ if not return_dict:
+ output = (feature_maps,)
+ if output_hidden_states:
+ output += (outputs.hidden_states,)
+ return output
+
+ return BackboneOutput(
+ feature_maps=feature_maps,
+ hidden_states=outputs.hidden_states if output_hidden_states else None,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = ["DinatForImageClassification", "DinatModel", "DinatPreTrainedModel", "DinatBackbone"]
diff --git a/third_party/transformers/src/transformers/models/dit/__init__.py b/third_party/transformers/src/transformers/models/dit/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/third_party/transformers/src/transformers/models/dit/convert_dit_unilm_to_pytorch.py b/third_party/transformers/src/transformers/models/dit/convert_dit_unilm_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..c627d60828705f5ad5fcecf02c506d1dad3afcb9
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/dit/convert_dit_unilm_to_pytorch.py
@@ -0,0 +1,221 @@
+# Copyright 2022 The HuggingFace Inc. team.
+#
+# 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.
+"""Convert DiT checkpoints from the unilm repository."""
+
+import argparse
+import json
+from io import BytesIO
+from pathlib import Path
+
+import httpx
+import torch
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+from transformers import BeitConfig, BeitForImageClassification, BeitForMaskedImageModeling, BeitImageProcessor
+from transformers.image_utils import PILImageResampling
+from transformers.utils import logging
+
+
+logging.set_verbosity_info()
+logger = logging.get_logger(__name__)
+
+
+# here we list all keys to be renamed (original name on the left, our name on the right)
+def create_rename_keys(config, has_lm_head=False, is_semantic=False):
+ prefix = "backbone." if is_semantic else ""
+
+ rename_keys = []
+ for i in range(config.num_hidden_layers):
+ # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
+ rename_keys.append((f"{prefix}blocks.{i}.norm1.weight", f"beit.encoder.layer.{i}.layernorm_before.weight"))
+ rename_keys.append((f"{prefix}blocks.{i}.norm1.bias", f"beit.encoder.layer.{i}.layernorm_before.bias"))
+ rename_keys.append(
+ (f"{prefix}blocks.{i}.attn.proj.weight", f"beit.encoder.layer.{i}.attention.output.dense.weight")
+ )
+ rename_keys.append(
+ (f"{prefix}blocks.{i}.attn.proj.bias", f"beit.encoder.layer.{i}.attention.output.dense.bias")
+ )
+ rename_keys.append((f"{prefix}blocks.{i}.norm2.weight", f"beit.encoder.layer.{i}.layernorm_after.weight"))
+ rename_keys.append((f"{prefix}blocks.{i}.norm2.bias", f"beit.encoder.layer.{i}.layernorm_after.bias"))
+ rename_keys.append((f"{prefix}blocks.{i}.mlp.fc1.weight", f"beit.encoder.layer.{i}.intermediate.dense.weight"))
+ rename_keys.append((f"{prefix}blocks.{i}.mlp.fc1.bias", f"beit.encoder.layer.{i}.intermediate.dense.bias"))
+ rename_keys.append((f"{prefix}blocks.{i}.mlp.fc2.weight", f"beit.encoder.layer.{i}.output.dense.weight"))
+ rename_keys.append((f"{prefix}blocks.{i}.mlp.fc2.bias", f"beit.encoder.layer.{i}.output.dense.bias"))
+
+ # projection layer + position embeddings
+ rename_keys.extend(
+ [
+ (f"{prefix}cls_token", "beit.embeddings.cls_token"),
+ (f"{prefix}patch_embed.proj.weight", "beit.embeddings.patch_embeddings.projection.weight"),
+ (f"{prefix}patch_embed.proj.bias", "beit.embeddings.patch_embeddings.projection.bias"),
+ (f"{prefix}pos_embed", "beit.embeddings.position_embeddings"),
+ ]
+ )
+
+ if has_lm_head:
+ # mask token + layernorm
+ rename_keys.extend(
+ [
+ ("mask_token", "beit.embeddings.mask_token"),
+ ("norm.weight", "layernorm.weight"),
+ ("norm.bias", "layernorm.bias"),
+ ]
+ )
+ else:
+ # layernorm + classification head
+ rename_keys.extend(
+ [
+ ("fc_norm.weight", "beit.pooler.layernorm.weight"),
+ ("fc_norm.bias", "beit.pooler.layernorm.bias"),
+ ("head.weight", "classifier.weight"),
+ ("head.bias", "classifier.bias"),
+ ]
+ )
+
+ return rename_keys
+
+
+# we split up the matrix of each encoder layer into queries, keys and values
+def read_in_q_k_v(state_dict, config, has_lm_head=False, is_semantic=False):
+ for i in range(config.num_hidden_layers):
+ prefix = "backbone." if is_semantic else ""
+ # queries, keys and values
+ in_proj_weight = state_dict.pop(f"{prefix}blocks.{i}.attn.qkv.weight")
+ q_bias = state_dict.pop(f"{prefix}blocks.{i}.attn.q_bias")
+ v_bias = state_dict.pop(f"{prefix}blocks.{i}.attn.v_bias")
+
+ state_dict[f"beit.encoder.layer.{i}.attention.attention.query.weight"] = in_proj_weight[
+ : config.hidden_size, :
+ ]
+ state_dict[f"beit.encoder.layer.{i}.attention.attention.query.bias"] = q_bias
+ state_dict[f"beit.encoder.layer.{i}.attention.attention.key.weight"] = in_proj_weight[
+ config.hidden_size : config.hidden_size * 2, :
+ ]
+ state_dict[f"beit.encoder.layer.{i}.attention.attention.value.weight"] = in_proj_weight[
+ -config.hidden_size :, :
+ ]
+ state_dict[f"beit.encoder.layer.{i}.attention.attention.value.bias"] = v_bias
+
+ # gamma_1 and gamma_2
+ # we call them lambda because otherwise they are renamed when using .from_pretrained
+ gamma_1 = state_dict.pop(f"{prefix}blocks.{i}.gamma_1")
+ gamma_2 = state_dict.pop(f"{prefix}blocks.{i}.gamma_2")
+
+ state_dict[f"beit.encoder.layer.{i}.lambda_1"] = gamma_1
+ state_dict[f"beit.encoder.layer.{i}.lambda_2"] = gamma_2
+
+
+def rename_key(dct, old, new):
+ val = dct.pop(old)
+ dct[new] = val
+
+
+# We will verify our results on an image of cute cats
+def prepare_img():
+ url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ with httpx.stream("GET", url) as response:
+ image = Image.open(BytesIO(response.read()))
+ return image
+
+
+@torch.no_grad()
+def convert_dit_checkpoint(checkpoint_url, pytorch_dump_folder_path, push_to_hub=False):
+ """
+ Copy/paste/tweak model's weights to our BEiT structure.
+ """
+
+ # define default BEiT configuration
+ has_lm_head = "rvlcdip" not in checkpoint_url
+ config = BeitConfig(use_absolute_position_embeddings=True, use_mask_token=has_lm_head)
+
+ # size of the architecture
+ if "large" in checkpoint_url or "dit-l" in checkpoint_url:
+ config.hidden_size = 1024
+ config.intermediate_size = 4096
+ config.num_hidden_layers = 24
+ config.num_attention_heads = 16
+
+ # labels
+ if "rvlcdip" in checkpoint_url:
+ config.num_labels = 16
+ repo_id = "huggingface/label-files"
+ filename = "rvlcdip-id2label.json"
+ id2label = json.load(open(hf_hub_download(repo_id, filename, repo_type="dataset"), "r"))
+ id2label = {int(k): v for k, v in id2label.items()}
+ config.id2label = id2label
+ config.label2id = {v: k for k, v in id2label.items()}
+
+ # load state_dict of original model, remove and rename some keys
+ state_dict = torch.hub.load_state_dict_from_url(checkpoint_url, map_location="cpu")["model"]
+
+ rename_keys = create_rename_keys(config, has_lm_head=has_lm_head)
+ for src, dest in rename_keys:
+ rename_key(state_dict, src, dest)
+ read_in_q_k_v(state_dict, config, has_lm_head=has_lm_head)
+
+ # load HuggingFace model
+ model = BeitForMaskedImageModeling(config) if has_lm_head else BeitForImageClassification(config)
+ model.eval()
+ model.load_state_dict(state_dict)
+
+ # Check outputs on an image
+ image_processor = BeitImageProcessor(
+ size=config.image_size, resample=PILImageResampling.BILINEAR, do_center_crop=False
+ )
+ image = prepare_img()
+
+ encoding = image_processor(images=image, return_tensors="pt")
+ pixel_values = encoding["pixel_values"]
+
+ outputs = model(pixel_values)
+ logits = outputs.logits
+
+ # verify logits
+ expected_shape = [1, 16] if "rvlcdip" in checkpoint_url else [1, 196, 8192]
+ assert logits.shape == torch.Size(expected_shape), "Shape of logits not as expected"
+
+ Path(pytorch_dump_folder_path).mkdir(exist_ok=True)
+ print(f"Saving model to {pytorch_dump_folder_path}")
+ model.save_pretrained(pytorch_dump_folder_path)
+ print(f"Saving image processor to {pytorch_dump_folder_path}")
+ image_processor.save_pretrained(pytorch_dump_folder_path)
+
+ if push_to_hub:
+ if has_lm_head:
+ model_name = "dit-base" if "base" in checkpoint_url else "dit-large"
+ else:
+ model_name = "dit-base-finetuned-rvlcdip" if "dit-b" in checkpoint_url else "dit-large-finetuned-rvlcdip"
+ image_processor.push_to_hub(repo_id=f"nielsr/{model_name}")
+ model.push_to_hub(repo_id=f"nielsr/{model_name}")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+
+ parser.add_argument(
+ "--checkpoint_url",
+ default="https://layoutlm.blob.core.windows.net/dit/dit-pts/dit-base-224-p16-500k-62d53a.pth",
+ type=str,
+ help="URL to the original PyTorch checkpoint (.pth file).",
+ )
+ parser.add_argument(
+ "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model."
+ )
+ parser.add_argument(
+ "--push_to_hub",
+ action="store_true",
+ )
+ args = parser.parse_args()
+ convert_dit_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path, args.push_to_hub)
diff --git a/third_party/transformers/src/transformers/models/electra/__init__.py b/third_party/transformers/src/transformers/models/electra/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..200b5e4106a0b200a9615a63f44b0def7c7e0650
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/electra/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from ..bert.tokenization_bert import BertTokenizer as ElectraTokenizer
+ from .configuration_electra import *
+ from .modeling_electra import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/electra/configuration_electra.py b/third_party/transformers/src/transformers/models/electra/configuration_electra.py
new file mode 100644
index 0000000000000000000000000000000000000000..19a332639deff6c085097f199b3d521c75d6cfd9
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/electra/configuration_electra.py
@@ -0,0 +1,89 @@
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+"""ELECTRA model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="google/electra-small-discriminator")
+@strict
+class ElectraConfig(PreTrainedConfig):
+ r"""
+ summary_type (`str`, *optional*, defaults to `"first"`):
+ Argument used when doing sequence summary. Used in the sequence classification and multiple choice models.
+ Has to be one of the following options:
+ - `"last"`: Take the last token hidden state (like XLNet).
+ - `"first"`: Take the first token hidden state (like BERT).
+ - `"mean"`: Take the mean of all tokens hidden states.
+ - `"cls_index"`: Supply a Tensor of classification token position (like GPT/GPT-2).
+ - `"attn"`: Not implemented now, use multi-head attention.
+ summary_use_proj (`bool`, *optional*, defaults to `True`):
+ Argument used when doing sequence summary. Used in the sequence classification and multiple choice models.
+ Whether or not to add a projection after the vector extraction.
+ summary_activation (`str`, *optional*):
+ Argument used when doing sequence summary. Used in the sequence classification and multiple choice models.
+ Pass `"gelu"` for a gelu activation to the output, any other value will result in no activation.
+ summary_last_dropout (`float`, *optional*, defaults to 0.0):
+ Argument used when doing sequence summary. Used in the sequence classification and multiple choice models.
+ The dropout ratio to be used after the projection and activation.
+
+ Examples:
+
+ ```python
+ >>> from transformers import ElectraConfig, ElectraModel
+
+ >>> # Initializing a ELECTRA electra-base-uncased style configuration
+ >>> configuration = ElectraConfig()
+
+ >>> # Initializing a model (with random weights) from the electra-base-uncased style configuration
+ >>> model = ElectraModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "electra"
+
+ vocab_size: int = 30522
+ embedding_size: int = 128
+ hidden_size: int = 256
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 4
+ intermediate_size: int = 1024
+ hidden_act: str = "gelu"
+ hidden_dropout_prob: float | int = 0.1
+ attention_probs_dropout_prob: float | int = 0.1
+ max_position_embeddings: int = 512
+ type_vocab_size: int = 2
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-12
+ summary_type: str = "first"
+ summary_use_proj: bool = True
+ summary_activation: str = "gelu"
+ summary_last_dropout: float | int = 0.1
+ pad_token_id: int | None = 0
+ use_cache: bool = True
+ classifier_dropout: float | int | None = None
+ is_decoder: bool = False
+ add_cross_attention: bool = False
+ bos_token_id: int | None = None
+ eos_token_id: int | list[int] | None = None
+ tie_word_embeddings: bool = True
+
+
+__all__ = ["ElectraConfig"]
diff --git a/third_party/transformers/src/transformers/models/electra/convert_electra_original_tf_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/electra/convert_electra_original_tf_checkpoint_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..4bb62f6858bad0b1a0b3b13e176759641828eb88
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/electra/convert_electra_original_tf_checkpoint_to_pytorch.py
@@ -0,0 +1,162 @@
+# Copyright 2018 The HuggingFace Inc. team.
+#
+# 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.
+"""Convert ELECTRA checkpoint."""
+
+import argparse
+import os
+
+import torch
+
+from transformers import ElectraConfig, ElectraForMaskedLM, ElectraForPreTraining
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+logging.set_verbosity_info()
+
+
+def load_tf_weights_in_electra(model, config, tf_checkpoint_path, discriminator_or_generator="discriminator"):
+ """Load tf checkpoints in a pytorch model."""
+ try:
+ import re
+
+ import numpy as np
+ import tensorflow as tf
+ except ImportError:
+ logger.error(
+ "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
+ "https://www.tensorflow.org/install/ for installation instructions."
+ )
+ raise
+ tf_path = os.path.abspath(tf_checkpoint_path)
+ logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
+ # Load weights from TF model
+ init_vars = tf.train.list_variables(tf_path)
+ names = []
+ arrays = []
+ for name, shape in init_vars:
+ logger.info(f"Loading TF weight {name} with shape {shape}")
+ array = tf.train.load_variable(tf_path, name)
+ names.append(name)
+ arrays.append(array)
+ for name, array in zip(names, arrays):
+ original_name: str = name
+
+ try:
+ if isinstance(model, ElectraForMaskedLM):
+ name = name.replace("electra/embeddings/", "generator/embeddings/")
+
+ if discriminator_or_generator == "generator":
+ name = name.replace("electra/", "discriminator/")
+ name = name.replace("generator/", "electra/")
+
+ name = name.replace("dense_1", "dense_prediction")
+ name = name.replace("generator_predictions/output_bias", "generator_lm_head/bias")
+
+ name = name.split("/")
+ # print(original_name, name)
+ # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
+ # which are not required for using pretrained model
+ if any(n in ["global_step", "temperature"] for n in name):
+ logger.info(f"Skipping {original_name}")
+ continue
+ pointer = model
+ for m_name in name:
+ if re.fullmatch(r"[A-Za-z]+_\d+", m_name):
+ scope_names = re.split(r"_(\d+)", m_name)
+ else:
+ scope_names = [m_name]
+ if scope_names[0] == "kernel" or scope_names[0] == "gamma":
+ pointer = getattr(pointer, "weight")
+ elif scope_names[0] == "output_bias" or scope_names[0] == "beta":
+ pointer = getattr(pointer, "bias")
+ elif scope_names[0] == "output_weights":
+ pointer = getattr(pointer, "weight")
+ elif scope_names[0] == "squad":
+ pointer = getattr(pointer, "classifier")
+ else:
+ pointer = getattr(pointer, scope_names[0])
+ if len(scope_names) >= 2:
+ num = int(scope_names[1])
+ pointer = pointer[num]
+ if m_name.endswith("_embeddings"):
+ pointer = getattr(pointer, "weight")
+ elif m_name == "kernel":
+ array = np.transpose(array)
+ try:
+ if pointer.shape != array.shape:
+ raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched")
+ except ValueError as e:
+ e.args += (pointer.shape, array.shape)
+ raise
+ print(f"Initialize PyTorch weight {name}", original_name)
+ pointer.data = torch.from_numpy(array)
+ except AttributeError as e:
+ print(f"Skipping {original_name}", name, e)
+ continue
+ return model
+
+
+def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, config_file, pytorch_dump_path, discriminator_or_generator):
+ # Initialise PyTorch model
+ config = ElectraConfig.from_json_file(config_file)
+ print(f"Building PyTorch model from configuration: {config}")
+
+ if discriminator_or_generator == "discriminator":
+ model = ElectraForPreTraining(config)
+ elif discriminator_or_generator == "generator":
+ model = ElectraForMaskedLM(config)
+ else:
+ raise ValueError("The discriminator_or_generator argument should be either 'discriminator' or 'generator'")
+
+ # Load weights from tf checkpoint
+ load_tf_weights_in_electra(
+ model, config, tf_checkpoint_path, discriminator_or_generator=discriminator_or_generator
+ )
+
+ # Save pytorch-model
+ print(f"Save PyTorch model to {pytorch_dump_path}")
+ torch.save(model.state_dict(), pytorch_dump_path)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ # Required parameters
+ parser.add_argument(
+ "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
+ )
+ parser.add_argument(
+ "--config_file",
+ default=None,
+ type=str,
+ required=True,
+ help="The config json file corresponding to the pre-trained model. \nThis specifies the model architecture.",
+ )
+ parser.add_argument(
+ "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
+ )
+ parser.add_argument(
+ "--discriminator_or_generator",
+ default=None,
+ type=str,
+ required=True,
+ help=(
+ "Whether to export the generator or the discriminator. Should be a string, either 'discriminator' or "
+ "'generator'."
+ ),
+ )
+ args = parser.parse_args()
+ convert_tf_checkpoint_to_pytorch(
+ args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.discriminator_or_generator
+ )
diff --git a/third_party/transformers/src/transformers/models/electra/modeling_electra.py b/third_party/transformers/src/transformers/models/electra/modeling_electra.py
new file mode 100644
index 0000000000000000000000000000000000000000..c69ad30c8746a382f9f1497a94c733a3e6b2ea3b
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/electra/modeling_electra.py
@@ -0,0 +1,1378 @@
+# Copyright 2019 The Google AI Language Team Authors and The HuggingFace Inc. team.
+#
+# 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 ELECTRA model."""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ... import initialization as init
+from ...activations import ACT2FN, get_activation
+from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_bidirectional_mask, create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutputWithCrossAttentions,
+ BaseModelOutputWithPastAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions,
+ MaskedLMOutput,
+ MultipleChoiceModelOutput,
+ QuestionAnsweringModelOutput,
+ SequenceClassifierOutput,
+ TokenClassifierOutput,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...pytorch_utils import apply_chunking_to_forward
+from ...utils import (
+ ModelOutput,
+ TransformersKwargs,
+ auto_docstring,
+ logging,
+)
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_electra import ElectraConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class ElectraEmbeddings(nn.Module):
+ """Construct the embeddings from word, position and token_type embeddings."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.embedding_size, padding_idx=config.pad_token_id)
+ self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.embedding_size)
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.embedding_size)
+
+ self.LayerNorm = nn.LayerNorm(config.embedding_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+ self.register_buffer(
+ "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+ )
+ self.register_buffer(
+ "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
+ )
+
+ # Copied from transformers.models.bert.modeling_bert.BertEmbeddings.forward
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values_length: int = 0,
+ ) -> torch.Tensor:
+ if input_ids is not None:
+ input_shape = input_ids.size()
+ else:
+ input_shape = inputs_embeds.size()[:-1]
+
+ batch_size, seq_length = input_shape
+
+ if position_ids is None:
+ position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
+
+ # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
+ # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
+ # issue #5664
+ if token_type_ids is None:
+ if hasattr(self, "token_type_ids"):
+ # NOTE: We assume either pos ids to have bsz == 1 (broadcastable) or bsz == effective bsz (input_shape[0])
+ buffered_token_type_ids = self.token_type_ids.expand(position_ids.shape[0], -1)
+ buffered_token_type_ids = torch.gather(buffered_token_type_ids, dim=1, index=position_ids)
+ token_type_ids = buffered_token_type_ids.expand(batch_size, seq_length)
+ else:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.word_embeddings(input_ids)
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
+ embeddings = inputs_embeds + token_type_embeddings
+
+ position_embeddings = self.position_embeddings(position_ids)
+ embeddings = embeddings + position_embeddings
+
+ embeddings = self.LayerNorm(embeddings)
+ embeddings = self.dropout(embeddings)
+ return embeddings
+
+
+# Copied from transformers.models.bert.modeling_bert.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float | None = None,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ if scaling is None:
+ scaling = query.size(-1) ** -0.5
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfAttention with Bert->Electra
+class ElectraSelfAttention(nn.Module):
+ def __init__(self, config, is_causal=False, layer_idx=None):
+ super().__init__()
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+ raise ValueError(
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+ f"heads ({config.num_attention_heads})"
+ )
+ self.config = config
+
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+ self.scaling = self.attention_head_size**-0.5
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ self.is_decoder = config.is_decoder
+ self.is_causal = is_causal
+ self.layer_idx = layer_idx
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.attention_head_size)
+
+ # get all proj
+ query_layer = self.query(hidden_states).view(*hidden_shape).transpose(1, 2)
+ key_layer = self.key(hidden_states).view(*hidden_shape).transpose(1, 2)
+ value_layer = self.value(hidden_states).view(*hidden_shape).transpose(1, 2)
+
+ if past_key_values is not None:
+ # decoder-only bert can have a simple dynamic cache for example
+ current_past_key_values = past_key_values
+ if isinstance(past_key_values, EncoderDecoderCache):
+ current_past_key_values = past_key_values.self_attention_cache
+
+ # save all key/value_layer to cache to be re-used for fast auto-regressive generation
+ key_layer, value_layer = current_past_key_values.update(key_layer, value_layer, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_layer,
+ key_layer,
+ value_layer,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout.p,
+ scaling=self.scaling,
+ **kwargs,
+ )
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.bert.modeling_bert.BertCrossAttention with Bert->Electra
+class ElectraCrossAttention(nn.Module):
+ def __init__(self, config, is_causal=False, layer_idx=None):
+ super().__init__()
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+ raise ValueError(
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+ f"heads ({config.num_attention_heads})"
+ )
+ self.config = config
+
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+ self.scaling = self.attention_head_size**-0.5
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ self.is_causal = is_causal
+ self.layer_idx = layer_idx
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ past_key_values: EncoderDecoderCache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor]:
+ # determine input shapes
+ bsz, tgt_len = hidden_states.shape[:-1]
+ src_len = encoder_hidden_states.shape[1]
+
+ q_input_shape = (bsz, tgt_len, -1, self.attention_head_size)
+ kv_input_shape = (bsz, src_len, -1, self.attention_head_size)
+
+ # get query proj
+ query_layer = self.query(hidden_states).view(*q_input_shape).transpose(1, 2)
+
+ is_updated = past_key_values.is_updated.get(self.layer_idx) if past_key_values is not None else False
+ if past_key_values is not None and is_updated:
+ # reuse k,v, cross_attentions
+ key_layer = past_key_values.cross_attention_cache.layers[self.layer_idx].keys
+ value_layer = past_key_values.cross_attention_cache.layers[self.layer_idx].values
+ else:
+ key_layer = self.key(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2)
+ value_layer = self.value(encoder_hidden_states).view(*kv_input_shape).transpose(1, 2)
+
+ if past_key_values is not None:
+ # save all states to the cache
+ key_layer, value_layer = past_key_values.cross_attention_cache.update(
+ key_layer, value_layer, self.layer_idx
+ )
+ # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls
+ past_key_values.is_updated[self.layer_idx] = True
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_layer,
+ key_layer,
+ value_layer,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout.p,
+ scaling=self.scaling,
+ **kwargs,
+ )
+ attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfOutput
+class ElectraSelfOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertAttention with Bert->Electra,BERT->ELECTRA
+class ElectraAttention(nn.Module):
+ def __init__(self, config, is_causal=False, layer_idx=None, is_cross_attention=False):
+ super().__init__()
+ self.is_cross_attention = is_cross_attention
+ attention_class = ElectraCrossAttention if is_cross_attention else ElectraSelfAttention
+ self.self = attention_class(config, is_causal=is_causal, layer_idx=layer_idx)
+ self.output = ElectraSelfOutput(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor]:
+ attention_mask = attention_mask if not self.is_cross_attention else encoder_attention_mask
+ attention_output, attn_weights = self.self(
+ hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ attention_output = self.output(attention_output, hidden_states)
+ return attention_output, attn_weights
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate
+class ElectraIntermediate(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertOutput
+class ElectraOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertLayer with Bert->Electra
+class ElectraLayer(GradientCheckpointingLayer):
+ def __init__(self, config, layer_idx=None):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+ self.attention = ElectraAttention(config, is_causal=config.is_decoder, layer_idx=layer_idx)
+ self.is_decoder = config.is_decoder
+ self.add_cross_attention = config.add_cross_attention
+ if self.add_cross_attention:
+ if not self.is_decoder:
+ raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
+ self.crossattention = ElectraAttention(
+ config,
+ is_causal=False,
+ layer_idx=layer_idx,
+ is_cross_attention=True,
+ )
+ self.intermediate = ElectraIntermediate(config)
+ self.output = ElectraOutput(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ self_attention_output, _ = self.attention(
+ hidden_states,
+ attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ attention_output = self_attention_output
+
+ if self.is_decoder and encoder_hidden_states is not None:
+ if not hasattr(self, "crossattention"):
+ raise ValueError(
+ f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
+ " by setting `config.add_cross_attention=True`"
+ )
+
+ cross_attention_output, _ = self.crossattention(
+ self_attention_output,
+ None, # attention_mask
+ encoder_hidden_states,
+ encoder_attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ attention_output = cross_attention_output
+
+ layer_output = apply_chunking_to_forward(
+ self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+ )
+ return layer_output
+
+ def feed_forward_chunk(self, attention_output):
+ intermediate_output = self.intermediate(attention_output)
+ layer_output = self.output(intermediate_output, attention_output)
+ return layer_output
+
+
+# Copied from transformers.models.bert.modeling_bert.BertEncoder with Bert->Electra
+class ElectraEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList([ElectraLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | BaseModelOutputWithPastAndCrossAttentions:
+ for i, layer_module in enumerate(self.layer):
+ hidden_states = layer_module(
+ hidden_states,
+ attention_mask,
+ encoder_hidden_states, # as a positional argument for gradient checkpointing
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values if use_cache else None,
+ )
+
+
+class ElectraDiscriminatorPredictions(nn.Module):
+ """Prediction module for the discriminator, made up of two dense layers."""
+
+ def __init__(self, config):
+ super().__init__()
+
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.activation = get_activation(config.hidden_act)
+ self.dense_prediction = nn.Linear(config.hidden_size, 1)
+ self.config = config
+
+ def forward(self, discriminator_hidden_states):
+ hidden_states = self.dense(discriminator_hidden_states)
+ hidden_states = self.activation(hidden_states)
+ logits = self.dense_prediction(hidden_states).squeeze(-1)
+
+ return logits
+
+
+class ElectraGeneratorPredictions(nn.Module):
+ """Prediction module for the generator, made up of two dense layers."""
+
+ def __init__(self, config):
+ super().__init__()
+
+ self.activation = get_activation("gelu")
+ self.LayerNorm = nn.LayerNorm(config.embedding_size, eps=config.layer_norm_eps)
+ self.dense = nn.Linear(config.hidden_size, config.embedding_size)
+
+ def forward(self, generator_hidden_states):
+ hidden_states = self.dense(generator_hidden_states)
+ hidden_states = self.activation(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states)
+
+ return hidden_states
+
+
+@auto_docstring
+class ElectraPreTrainedModel(PreTrainedModel):
+ config_class = ElectraConfig
+ base_model_prefix = "electra"
+ supports_gradient_checkpointing = True
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": ElectraLayer,
+ "attentions": ElectraSelfAttention,
+ "cross_attentions": ElectraCrossAttention,
+ }
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, ElectraEmbeddings):
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+ init.zeros_(module.token_type_ids)
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Output type of [`ElectraForPreTraining`].
+ """
+)
+class ElectraForPreTrainingOutput(ModelOutput):
+ r"""
+ loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
+ Total loss of the ELECTRA objective.
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
+ Prediction scores of the head (scores for each token before SoftMax).
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring
+class ElectraModel(ElectraPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.embeddings = ElectraEmbeddings(config)
+
+ if config.embedding_size != config.hidden_size:
+ self.embeddings_project = nn.Linear(config.embedding_size, config.hidden_size)
+
+ self.encoder = ElectraEncoder(config)
+ self.config = config
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embeddings.word_embeddings
+
+ def set_input_embeddings(self, value):
+ self.embeddings.word_embeddings = value
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ past_key_values: list[torch.FloatTensor] | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | BaseModelOutputWithCrossAttentions:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if self.config.is_decoder:
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+ else:
+ use_cache = False
+
+ if use_cache and past_key_values is None:
+ past_key_values = (
+ EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))
+ if encoder_hidden_states is not None or self.config.is_encoder_decoder
+ else DynamicCache(config=self.config)
+ )
+
+ past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+
+ embedding_output = self.embeddings(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ past_key_values_length=past_key_values_length,
+ )
+ if hasattr(self, "embeddings_project"):
+ embedding_output = self.embeddings_project(embedding_output)
+
+ attention_mask, encoder_attention_mask = self._create_attention_masks(
+ attention_mask=attention_mask,
+ encoder_attention_mask=encoder_attention_mask,
+ embedding_output=embedding_output,
+ encoder_hidden_states=encoder_hidden_states,
+ past_key_values=past_key_values,
+ )
+
+ encoder_outputs = self.encoder(
+ embedding_output,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=encoder_outputs.last_hidden_state,
+ past_key_values=encoder_outputs.past_key_values,
+ )
+
+ # Copied from transformers.models.bert.modeling_bert.BertModel._create_attention_masks
+ def _create_attention_masks(
+ self,
+ attention_mask,
+ encoder_attention_mask,
+ embedding_output,
+ encoder_hidden_states,
+ past_key_values,
+ ):
+ if self.config.is_decoder:
+ attention_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=embedding_output,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ )
+ else:
+ attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=embedding_output,
+ attention_mask=attention_mask,
+ )
+
+ if encoder_attention_mask is not None:
+ encoder_attention_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=embedding_output,
+ attention_mask=encoder_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ )
+
+ return attention_mask, encoder_attention_mask
+
+
+class ElectraClassificationHead(nn.Module):
+ """Head for sentence-level classification tasks."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ classifier_dropout = (
+ config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
+ )
+ self.activation = get_activation("gelu")
+ self.dropout = nn.Dropout(classifier_dropout)
+ self.out_proj = nn.Linear(config.hidden_size, config.num_labels)
+
+ def forward(self, features, **kwargs):
+ x = features[:, 0, :] # take token (equiv. to [CLS])
+ x = self.dropout(x)
+ x = self.dense(x)
+ x = self.activation(x) # although BERT uses tanh here, it seems Electra authors used gelu here
+ x = self.dropout(x)
+ x = self.out_proj(x)
+ return x
+
+
+# Copied from transformers.models.xlm.modeling_xlm.XLMSequenceSummary with XLM->Electra
+class ElectraSequenceSummary(nn.Module):
+ r"""
+ Compute a single vector summary of a sequence hidden states.
+
+ Args:
+ config ([`ElectraConfig`]):
+ The config used by the model. Relevant arguments in the config class of the model are (refer to the actual
+ config class of your model for the default values it uses):
+
+ - **summary_type** (`str`) -- The method to use to make this summary. Accepted values are:
+
+ - `"last"` -- Take the last token hidden state (like XLNet)
+ - `"first"` -- Take the first token hidden state (like Bert)
+ - `"mean"` -- Take the mean of all tokens hidden states
+ - `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2)
+ - `"attn"` -- Not implemented now, use multi-head attention
+
+ - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction.
+ - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes
+ (otherwise to `config.hidden_size`).
+ - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output,
+ another string or `None` will add no activation.
+ - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation.
+ - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation.
+ """
+
+ def __init__(self, config: ElectraConfig):
+ super().__init__()
+
+ self.summary_type = getattr(config, "summary_type", "last")
+ if self.summary_type == "attn":
+ # We should use a standard multi-head attention module with absolute positional embedding for that.
+ # Cf. https://github.com/zihangdai/xlnet/blob/master/modeling.py#L253-L276
+ # We can probably just use the multi-head attention module of PyTorch >=1.1.0
+ raise NotImplementedError
+
+ self.summary = nn.Identity()
+ if hasattr(config, "summary_use_proj") and config.summary_use_proj:
+ if hasattr(config, "summary_proj_to_labels") and config.summary_proj_to_labels and config.num_labels > 0:
+ num_classes = config.num_labels
+ else:
+ num_classes = config.hidden_size
+ self.summary = nn.Linear(config.hidden_size, num_classes)
+
+ activation_string = getattr(config, "summary_activation", None)
+ self.activation: Callable = get_activation(activation_string) if activation_string else nn.Identity()
+
+ self.first_dropout = nn.Identity()
+ if hasattr(config, "summary_first_dropout") and config.summary_first_dropout > 0:
+ self.first_dropout = nn.Dropout(config.summary_first_dropout)
+
+ self.last_dropout = nn.Identity()
+ if hasattr(config, "summary_last_dropout") and config.summary_last_dropout > 0:
+ self.last_dropout = nn.Dropout(config.summary_last_dropout)
+
+ def forward(
+ self, hidden_states: torch.FloatTensor, cls_index: torch.LongTensor | None = None
+ ) -> torch.FloatTensor:
+ """
+ Compute a single vector summary of a sequence hidden states.
+
+ Args:
+ hidden_states (`torch.FloatTensor` of shape `[batch_size, seq_len, hidden_size]`):
+ The hidden states of the last layer.
+ cls_index (`torch.LongTensor` of shape `[batch_size]` or `[batch_size, ...]` where ... are optional leading dimensions of `hidden_states`, *optional*):
+ Used if `summary_type == "cls_index"` and takes the last token of the sequence as classification token.
+
+ Returns:
+ `torch.FloatTensor`: The summary of the sequence hidden states.
+ """
+ if self.summary_type == "last":
+ output = hidden_states[:, -1]
+ elif self.summary_type == "first":
+ output = hidden_states[:, 0]
+ elif self.summary_type == "mean":
+ output = hidden_states.mean(dim=1)
+ elif self.summary_type == "cls_index":
+ if cls_index is None:
+ cls_index = torch.full_like(
+ hidden_states[..., :1, :],
+ hidden_states.shape[-2] - 1,
+ dtype=torch.long,
+ )
+ else:
+ cls_index = cls_index.unsqueeze(-1).unsqueeze(-1)
+ cls_index = cls_index.expand((-1,) * (cls_index.dim() - 1) + (hidden_states.size(-1),))
+ # shape of cls_index: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states
+ output = hidden_states.gather(-2, cls_index).squeeze(-2) # shape (bsz, XX, hidden_size)
+ elif self.summary_type == "attn":
+ raise NotImplementedError
+
+ output = self.first_dropout(output)
+ output = self.summary(output)
+ output = self.activation(output)
+ output = self.last_dropout(output)
+
+ return output
+
+
+@auto_docstring(
+ custom_intro="""
+ ELECTRA Model transformer with a sequence classification/regression head on top (a linear layer on top of the
+ pooled output) e.g. for GLUE tasks.
+ """
+)
+class ElectraForSequenceClassification(ElectraPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.config = config
+ self.electra = ElectraModel(config)
+ self.classifier = ElectraClassificationHead(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | SequenceClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ discriminator_hidden_states = self.electra(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ return_dict=True,
+ **kwargs,
+ )
+
+ sequence_output = discriminator_hidden_states[0]
+ logits = self.classifier(sequence_output)
+
+ loss = None
+ if labels is not None:
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(logits, labels)
+
+ return SequenceClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=discriminator_hidden_states.hidden_states,
+ attentions=discriminator_hidden_states.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Electra model with a binary classification head on top as used during pretraining for identifying generated tokens.
+
+ It is recommended to load the discriminator checkpoint into that model.
+ """
+)
+class ElectraForPreTraining(ElectraPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.electra = ElectraModel(config)
+ self.discriminator_predictions = ElectraDiscriminatorPredictions(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | ElectraForPreTrainingOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the ELECTRA loss. Input should be a sequence of tokens (see `input_ids` docstring)
+ Indices should be in `[0, 1]`:
+
+ - 0 indicates the token is an original token,
+ - 1 indicates the token was replaced.
+
+ Examples:
+
+ ```python
+ >>> from transformers import ElectraForPreTraining, AutoTokenizer
+ >>> import torch
+
+ >>> discriminator = ElectraForPreTraining.from_pretrained("google/electra-base-discriminator")
+ >>> tokenizer = AutoTokenizer.from_pretrained("google/electra-base-discriminator")
+
+ >>> sentence = "The quick brown fox jumps over the lazy dog"
+ >>> fake_sentence = "The quick brown fox fake over the lazy dog"
+
+ >>> fake_tokens = tokenizer.tokenize(fake_sentence, add_special_tokens=True)
+ >>> fake_inputs = tokenizer.encode(fake_sentence, return_tensors="pt")
+ >>> discriminator_outputs = discriminator(fake_inputs)
+ >>> predictions = torch.round((torch.sign(discriminator_outputs[0]) + 1) / 2)
+
+ >>> fake_tokens
+ ['[CLS]', 'the', 'quick', 'brown', 'fox', 'fake', 'over', 'the', 'lazy', 'dog', '[SEP]']
+
+ >>> predictions.squeeze().tolist()
+ [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0]
+ ```"""
+ discriminator_hidden_states = self.electra(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ return_dict=True,
+ **kwargs,
+ )
+ discriminator_sequence_output = discriminator_hidden_states[0]
+
+ logits = self.discriminator_predictions(discriminator_sequence_output)
+
+ loss = None
+ if labels is not None:
+ loss_fct = nn.BCEWithLogitsLoss()
+ if attention_mask is not None:
+ active_loss = attention_mask.view(-1, discriminator_sequence_output.shape[1]) == 1
+ active_logits = logits.view(-1, discriminator_sequence_output.shape[1])[active_loss]
+ active_labels = labels[active_loss]
+ loss = loss_fct(active_logits, active_labels.float())
+ else:
+ loss = loss_fct(logits.view(-1, discriminator_sequence_output.shape[1]), labels.float())
+
+ return ElectraForPreTrainingOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=discriminator_hidden_states.hidden_states,
+ attentions=discriminator_hidden_states.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Electra model with a language modeling head on top.
+
+ Even though both the discriminator and generator may be loaded into this model, the generator is the only model of
+ the two to have been trained for the masked language modeling task.
+ """
+)
+class ElectraForMaskedLM(ElectraPreTrainedModel):
+ _tied_weights_keys = {"generator_lm_head.weight": "electra.embeddings.word_embeddings.weight"}
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.electra = ElectraModel(config)
+ self.generator_predictions = ElectraGeneratorPredictions(config)
+
+ self.generator_lm_head = nn.Linear(config.embedding_size, config.vocab_size)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.generator_lm_head
+
+ def set_output_embeddings(self, word_embeddings):
+ self.generator_lm_head = word_embeddings
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | MaskedLMOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
+ config.vocab_size]` (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]`
+ """
+ generator_hidden_states = self.electra(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ return_dict=True,
+ **kwargs,
+ )
+ generator_sequence_output = generator_hidden_states[0]
+
+ prediction_scores = self.generator_predictions(generator_sequence_output)
+ prediction_scores = self.generator_lm_head(prediction_scores)
+
+ loss = None
+ # Masked language modeling softmax layer
+ if labels is not None:
+ loss_fct = nn.CrossEntropyLoss() # -100 index = padding token
+ loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+
+ return MaskedLMOutput(
+ loss=loss,
+ logits=prediction_scores,
+ hidden_states=generator_hidden_states.hidden_states,
+ attentions=generator_hidden_states.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Electra model with a token classification head on top.
+
+ Both the discriminator and generator may be loaded into this model.
+ """
+)
+class ElectraForTokenClassification(ElectraPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.electra = ElectraModel(config)
+ classifier_dropout = (
+ config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
+ )
+ self.dropout = nn.Dropout(classifier_dropout)
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | TokenClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
+ """
+ discriminator_hidden_states = self.electra(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ return_dict=True,
+ **kwargs,
+ )
+ discriminator_sequence_output = discriminator_hidden_states[0]
+
+ discriminator_sequence_output = self.dropout(discriminator_sequence_output)
+ logits = self.classifier(discriminator_sequence_output)
+
+ loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+ return TokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=discriminator_hidden_states.hidden_states,
+ attentions=discriminator_hidden_states.attentions,
+ )
+
+
+@auto_docstring
+class ElectraForQuestionAnswering(ElectraPreTrainedModel):
+ config_class = ElectraConfig
+ base_model_prefix = "electra"
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.electra = ElectraModel(config)
+ self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ start_positions: torch.Tensor | None = None,
+ end_positions: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | QuestionAnsweringModelOutput:
+ discriminator_hidden_states = self.electra(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ return_dict=True,
+ **kwargs,
+ )
+
+ sequence_output = discriminator_hidden_states[0]
+
+ logits = self.qa_outputs(sequence_output)
+ start_logits, end_logits = logits.split(1, dim=-1)
+ start_logits = start_logits.squeeze(-1).contiguous()
+ end_logits = end_logits.squeeze(-1).contiguous()
+
+ total_loss = None
+ if start_positions is not None and end_positions is not None:
+ # If we are on multi-GPU, split add a dimension
+ if len(start_positions.size()) > 1:
+ start_positions = start_positions.squeeze(-1)
+ if len(end_positions.size()) > 1:
+ end_positions = end_positions.squeeze(-1)
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
+ ignored_index = start_logits.size(1)
+ start_positions = start_positions.clamp(0, ignored_index)
+ end_positions = end_positions.clamp(0, ignored_index)
+
+ loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+ start_loss = loss_fct(start_logits, start_positions)
+ end_loss = loss_fct(end_logits, end_positions)
+ total_loss = (start_loss + end_loss) / 2
+
+ return QuestionAnsweringModelOutput(
+ loss=total_loss,
+ start_logits=start_logits,
+ end_logits=end_logits,
+ hidden_states=discriminator_hidden_states.hidden_states,
+ attentions=discriminator_hidden_states.attentions,
+ )
+
+
+@auto_docstring
+class ElectraForMultipleChoice(ElectraPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.electra = ElectraModel(config)
+ self.sequence_summary = ElectraSequenceSummary(config)
+ self.classifier = nn.Linear(config.hidden_size, 1)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | MultipleChoiceModelOutput:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ token_type_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*):
+ Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
+ 1]`:
+
+ - 0 corresponds to a *sentence A* token,
+ - 1 corresponds to a *sentence B* token.
+
+ [What are token type IDs?](../glossary#token-type-ids)
+ position_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*):
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+ config.max_position_embeddings - 1]`.
+
+ [What are position IDs?](../glossary#position-ids)
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_choices, 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.
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
+ num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
+ `input_ids` above)
+ """
+ num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
+
+ input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
+ attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
+ token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
+ position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+ inputs_embeds = (
+ inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
+ if inputs_embeds is not None
+ else None
+ )
+
+ discriminator_hidden_states = self.electra(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ return_dict=True,
+ **kwargs,
+ )
+
+ sequence_output = discriminator_hidden_states[0]
+
+ pooled_output = self.sequence_summary(sequence_output)
+ logits = self.classifier(pooled_output)
+ reshaped_logits = logits.view(-1, num_choices)
+
+ loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(reshaped_logits, labels)
+
+ return MultipleChoiceModelOutput(
+ loss=loss,
+ logits=reshaped_logits,
+ hidden_states=discriminator_hidden_states.hidden_states,
+ attentions=discriminator_hidden_states.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ ELECTRA Model with a `language modeling` head on top for CLM fine-tuning.
+ """
+)
+class ElectraForCausalLM(ElectraPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"generator_lm_head.weight": "electra.embeddings.word_embeddings.weight"}
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ if not config.is_decoder:
+ logger.warning("If you want to use `ElectraForCausalLM` as a standalone, add `is_decoder=True.`")
+
+ self.electra = ElectraModel(config)
+ self.generator_predictions = ElectraGeneratorPredictions(config)
+ self.generator_lm_head = nn.Linear(config.embedding_size, config.vocab_size)
+
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.generator_lm_head
+
+ def set_output_embeddings(self, new_embeddings):
+ self.generator_lm_head = new_embeddings
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ encoder_hidden_states: torch.Tensor | None = None,
+ encoder_attention_mask: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | CausalLMOutputWithCrossAttentions:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
+ `[-100, 0, ..., config.vocab_size]` (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]`
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, ElectraForCausalLM, ElectraConfig
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("google/electra-base-generator")
+ >>> config = ElectraConfig.from_pretrained("google/electra-base-generator")
+ >>> config.is_decoder = True
+ >>> model = ElectraForCausalLM.from_pretrained("google/electra-base-generator", config=config)
+
+ >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+ >>> outputs = model(**inputs)
+
+ >>> prediction_logits = outputs.logits
+ ```"""
+ if labels is not None:
+ use_cache = False
+
+ outputs: BaseModelOutputWithPastAndCrossAttentions = self.electra(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ return_dict=True,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.generator_lm_head(self.generator_predictions(hidden_states[:, slice_indices, :]))
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithCrossAttentions(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ cross_attentions=outputs.cross_attentions,
+ )
+
+
+__all__ = [
+ "ElectraForCausalLM",
+ "ElectraForMaskedLM",
+ "ElectraForMultipleChoice",
+ "ElectraForPreTraining",
+ "ElectraForQuestionAnswering",
+ "ElectraForSequenceClassification",
+ "ElectraForTokenClassification",
+ "ElectraModel",
+ "ElectraPreTrainedModel",
+]
diff --git a/third_party/transformers/src/transformers/models/emu3/__init__.py b/third_party/transformers/src/transformers/models/emu3/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d8555f58d1866451c38abb5559ef5bef9545f0b0
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/emu3/__init__.py
@@ -0,0 +1,29 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_emu3 import *
+ from .image_processing_emu3 import *
+ from .modeling_emu3 import *
+ from .processing_emu3 import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/emu3/configuration_emu3.py b/third_party/transformers/src/transformers/models/emu3/configuration_emu3.py
new file mode 100644
index 0000000000000000000000000000000000000000..0093891159e9ee14382a739008b80ad61b9409a0
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/emu3/configuration_emu3.py
@@ -0,0 +1,153 @@
+# Copyright 2024 HuggingFace Inc. team. All rights reserved.
+#
+#
+# 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.
+
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="Emu3-community/Emu3-Chat-hf")
+@strict
+class Emu3VQVAEConfig(PreTrainedConfig):
+ r"""
+ embed_dim (`int`, *optional*, defaults to 4):
+ Dimension of the quantized vector in codebook.
+ out_channels (`int`, *optional*, defaults to 3):
+ Output channel of decoder.
+ temporal_downsample_factor (`int`, *optional*, defaults to 4):
+ Temporal downsample factor.
+ base_channels (`int`, *optional*, defaults to 256):
+ Basic channel number of the intermediate blocks.
+ channel_multiplier (`list[int]`, *optional*, defaults to `[1, 2, 2, 4]`):
+ Channel scaling factor of the intermediate blocks.
+ num_res_blocks (`int`, *optional*, defaults to 2):
+ Residual block number in each stage.
+ attn_resolutions (`list[int]`, *optional*, defaults to `[3]`):
+ Stage indices to apply attention.
+
+ ```python
+ >>> from transformers import Emu3VQVAE, Emu3VQVAEConfig
+
+ >>> # Initializing a video VQ model of Emu3 configuration
+ >>> configuration = Emu3VQVAEConfig()
+
+ >>> # Initializing a model from the Emu3 VQ model style configuration
+ >>> model = Emu3VQVAE(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+ """
+
+ model_type = "emu3_vqgan"
+ base_config_key = "vq_config"
+
+ codebook_size: int = 32768
+ embed_dim: int = 4
+ latent_channels: int = 4
+ double_latent: bool = False
+ in_channels: int = 3
+ out_channels: int = 3
+ temporal_downsample_factor: int = 4
+ base_channels: int = 256
+ channel_multiplier: list[int] | tuple[int, ...] = (1, 2, 2, 4)
+ num_res_blocks: int = 2
+ attn_resolutions: list[int] | tuple[int, ...] = (3,)
+ hidden_size: int = 1024
+ num_attention_heads: int = 1
+ attention_dropout: float | int = 0.0
+
+
+@auto_docstring(checkpoint="Emu3-community/Emu3-Chat-hf")
+@strict
+class Emu3TextConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import Emu3Model, Emu3Config
+
+ >>> # Initializing a Emu3-community/Emu3-Chat-hf style configuration
+ >>> configuration = Emu3Config()
+
+ >>> # Initializing a model from the Emu3-community/Emu3-Chat-hf style configuration
+ >>> model = Emu3Model(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "emu3_text_model"
+ base_config_key = "text_config"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ default_theta = 1000000.0
+
+ vocab_size: int = 184622
+ hidden_size: int = 4096
+ intermediate_size: int = 14336
+ num_hidden_layers: int = 32
+ num_attention_heads: int = 32
+ num_key_value_heads: int | None = 8
+ hidden_act: str = "silu"
+ max_position_embeddings: int = 9216
+ rms_norm_eps: float = 1e-5
+ use_cache: bool = True
+ pad_token_id: int = 151643
+ bos_token_id: int = 151849
+ eos_token_id: int | list[int] | None = 151850
+ rope_parameters: RopeParameters | dict | None = None
+ mlp_bias = False
+ attention_bias = False
+ attention_dropout: float | int = 0.1
+ initializer_range: float = 0.02
+ tie_word_embeddings: bool = False
+
+
+@auto_docstring(checkpoint="Emu3-community/Emu3-Chat-hf")
+@strict
+class Emu3Config(PreTrainedConfig):
+ r"""
+ vocabulary_map (`dict`, *optional*):
+ A dictionary containing the vocabulary map from the tokenizer. Used to obtain tokens from the image inputs.
+ """
+
+ model_type = "emu3"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ sub_configs = {"text_config": Emu3TextConfig, "vq_config": Emu3VQVAEConfig}
+
+ vq_config: dict | Emu3VQVAEConfig | None = None
+ text_config: dict | Emu3TextConfig | None = None
+ vocabulary_map: dict[str, int] | None = None
+ tie_word_embeddings: bool = False
+
+ def __post_init__(self, **kwargs):
+ if self.vq_config is None:
+ self.vq_config = Emu3VQVAEConfig()
+ elif isinstance(self.vq_config, dict):
+ self.vq_config = Emu3VQVAEConfig(**self.vq_config)
+
+ if self.text_config is None:
+ self.text_config = Emu3TextConfig()
+ elif isinstance(self.text_config, dict):
+ self.text_config = Emu3TextConfig(**self.text_config)
+
+ self.image_token_id = self.vocabulary_map.get("") if self.vocabulary_map is not None else None
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["Emu3Config", "Emu3TextConfig", "Emu3VQVAEConfig"]
diff --git a/third_party/transformers/src/transformers/models/emu3/convert_emu3_weights_to_hf.py b/third_party/transformers/src/transformers/models/emu3/convert_emu3_weights_to_hf.py
new file mode 100644
index 0000000000000000000000000000000000000000..b86f92c865ad5fc8e562590f2de0d3c7b7c28179
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/emu3/convert_emu3_weights_to_hf.py
@@ -0,0 +1,444 @@
+# Copyright 2024 The Emu team, BAAI and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+import argparse
+import json
+import os
+import re
+from io import BytesIO
+
+import httpx
+import torch
+from PIL import Image
+
+from transformers import (
+ AutoModel,
+ AutoModelForCausalLM,
+ AutoTokenizer,
+ Emu3Config,
+ Emu3ForConditionalGeneration,
+ Emu3ImageProcessor,
+ Emu3Processor,
+ Emu3TextConfig,
+ GenerationConfig,
+)
+from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode
+
+
+"""
+Sample usage:
+
+```
+python src/transformers/models/emu3/convert_emu3_weights_to_hf.py \
+ --vq_model_id BAAI/Emu3-VisionTokenizer --llm_model_id BAAI/Emu3-Chat --output_dir /output/path
+```
+
+Thereafter, models can be loaded via:
+
+```py
+from transformers import Emu3ForConditionalGeneration, Emu3Processor
+
+model = Emu3ForConditionalGeneration.from_pretrained("/output/path")
+processor = Emu3Processor.from_pretrained("/output/path")
+```
+
+"""
+
+
+byte_encoder = bytes_to_unicode()
+CHAT_TEMPLATE = "{% for message in messages %}{% if message['role'] != 'system' %}{{ message['role'].upper() + ': '}}{% endif %}{# Render all images first #}{% for content in message['content'] | selectattr('type', 'equalto', 'image') %}{{ '' }}{% endfor %}{# Render all text next #}{% if message['role'] != 'assistant' %}{% for content in message['content'] | selectattr('type', 'equalto', 'text') %}{{ content['text'] + ' '}}{% endfor %}{% else %}{% for content in message['content'] | selectattr('type', 'equalto', 'text') %}{% generation %}{{ content['text'] + ' '}}{% endgeneration %}{% endfor %}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'ASSISTANT:' }}{% endif %}"
+
+
+# Tiktoken to HF conversion, thanks for Xenova
+def token_bytes_to_string(b):
+ return "".join([byte_encoder[ord(char)] for char in b.decode("latin-1")])
+
+
+# Adapted from https://github.com/openai/tiktoken/issues/60#issuecomment-1499977960
+def bpe(mergeable_ranks: dict[bytes, int], token: bytes, max_rank: int | None = None):
+ parts = [bytes([b]) for b in token]
+ while True:
+ min_idx = None
+ min_rank = None
+ for i, pair in enumerate(zip(parts[:-1], parts[1:])):
+ rank = mergeable_ranks.get(pair[0] + pair[1])
+ if rank is not None and (min_rank is None or rank < min_rank):
+ min_idx = i
+ min_rank = rank
+ if min_rank is None or (max_rank is not None and min_rank >= max_rank):
+ break
+ assert min_idx is not None
+ parts = parts[:min_idx] + [parts[min_idx] + parts[min_idx + 1]] + parts[min_idx + 2 :]
+ return parts
+
+
+def generate_vocab_and_merges(encoder):
+ mergeable_ranks = encoder._mergeable_ranks
+
+ merges = []
+ vocab = {}
+ for token, rank in mergeable_ranks.items():
+ vocab[token_bytes_to_string(token)] = rank
+
+ if len(token) == 1:
+ continue
+ merged = tuple(bpe(mergeable_ranks, token, max_rank=rank))
+ assert len(merged) == 2
+ merges.append(" ".join(map(token_bytes_to_string, merged)))
+
+ # Also add special tokens
+ vocab.update(encoder._special_tokens)
+ return vocab, merges
+
+
+def convert_tiktoken(tokenizer, output_dir):
+ encoder = tokenizer.tokenizer
+ vocab, merges = generate_vocab_and_merges(encoder)
+ added_tokens = [
+ {
+ "id": id,
+ "content": content,
+ "single_word": False,
+ "lstrip": False,
+ "rstrip": False,
+ "normalized": False,
+ "special": True,
+ }
+ for content, id in encoder._special_tokens.items()
+ if content != "<|extra_0|>"
+ ]
+
+ # https://huggingface.co/Xenova/gpt2/raw/main/tokenizer_config.json
+ tokenizer_config_template = {
+ "add_prefix_space": False,
+ "bos_token": "<|extra_203|>",
+ "clean_up_tokenization_spaces": False,
+ "eos_token": "<|extra_204|>",
+ "pad_token": "<|endoftext|>",
+ }
+ tokenizer_config_template.update({"tokenizer_class": "GPT2Tokenizer"})
+ tokenizer_config_template = dict(sorted(tokenizer_config_template.items(), key=lambda x: x[0]))
+
+ # add placeholder image token by taking one of the reserved tokens
+ reserved_token_id = vocab["<|extra_0|>"]
+ vocab[""] = reserved_token_id
+ del vocab["<|extra_0|>"]
+ added_tokens.append(
+ {
+ "id": reserved_token_id,
+ "content": "",
+ "single_word": False,
+ "lstrip": False,
+ "rstrip": False,
+ "normalized": False,
+ "special": True,
+ }
+ )
+
+ os.makedirs(output_dir, exist_ok=True)
+
+ pre_tokenizer = {
+ "type": "ByteLevel",
+ "add_prefix_space": False,
+ "trim_offsets": True,
+ "use_regex": True,
+ }
+
+ # https://huggingface.co/Xenova/gpt2/raw/main/tokenizer.json
+ tokenizer_template = {
+ "version": "1.0",
+ "truncation": None,
+ "padding": None,
+ "added_tokens": added_tokens,
+ "normalizer": None,
+ "pre_tokenizer": pre_tokenizer,
+ "post_processor": None,
+ "decoder": {
+ "type": "ByteLevel",
+ "add_prefix_space": True,
+ "trim_offsets": True,
+ "use_regex": True,
+ },
+ "model": {
+ "type": "BPE",
+ "dropout": None,
+ "unk_token": None,
+ "continuing_subword_prefix": "",
+ "end_of_word_suffix": "",
+ "fuse_unk": False,
+ "byte_fallback": False,
+ "vocab": vocab,
+ "merges": merges,
+ },
+ }
+
+ # Save to files
+ with open(os.path.join(output_dir, "vocab.json"), "w", encoding="utf-8") as fp:
+ json.dump(vocab, fp, indent=2, ensure_ascii=False)
+
+ with open(os.path.join(output_dir, "tokenizer.json"), "w", encoding="utf-8") as fp:
+ json.dump(tokenizer_template, fp, indent=2, ensure_ascii=False)
+
+ with open(os.path.join(output_dir, "tokenizer_config.json"), "w", encoding="utf-8") as fp:
+ json.dump(tokenizer_config_template, fp, indent=2, ensure_ascii=False)
+
+ with open(os.path.join(output_dir, "special_tokens_map.json"), "w", encoding="utf-8") as fp:
+ json.dump(
+ {
+ "bos_token": "<|extra_203|>",
+ "eos_token": "<|extra_204|>",
+ "pad_token": "<|endoftext|>",
+ },
+ fp,
+ indent=2,
+ ensure_ascii=False,
+ )
+
+ with open(os.path.join(output_dir, "merges.txt"), "w", encoding="utf-8") as fp:
+ fp.write("#version: 0.2\n")
+ fp.write("\n".join(merges))
+
+
+KEYS_TO_MODIFY_MAPPING = {
+ "^model": "model.text_model",
+ "^encoder": "model.vqmodel.encoder",
+ "^decoder": "model.vqmodel.decoder",
+ "^post_quant_conv": "model.vqmodel.post_quant_conv",
+ "^quant_conv": "model.vqmodel.quant_conv",
+ "^quantize": "model.vqmodel.quantize",
+ r"lm_head\.weight": "lm_head.weight",
+ # rename QKV proj for the VQ-VAE model because we use SiglipAttention
+ r"\.q\.": ".q_proj.",
+ r"\.k\.": ".k_proj.",
+ r"\.v\.": ".v_proj.",
+ r"\.proj_out\.": ".out_proj.",
+ # move the attention norms outside of attention modules
+ r"mid\.attn_1\.norm\.": "mid.attn_norm.",
+ r"attn\.0\.norm\.": "attn_norms.0.",
+ r"attn\.1\.norm\.": "attn_norms.1.",
+ r"attn\.2\.norm\.": "attn_norms.2.",
+ r"attn\.3\.norm\.": "attn_norms.3.",
+ # isolate down/mid/up into separate classes for readability
+ r"\.down\.": ".down_block.down.",
+ r"\.up\.": ".up_block.up.",
+ r"\.mid\.": ".middle_block.",
+}
+
+
+def convert_state_dict_to_hf(old_state_dict, new_state_dict):
+ for key, value in old_state_dict.items():
+ # convert conv layers in attn to linear
+ if (
+ any(key.endswith(name) for name in ["q.weight", "k.weight", "v.weight", "proj_out.weight"])
+ and value.ndim == 4
+ ):
+ value = value.squeeze()
+
+ for old_pattern, new_pattern in KEYS_TO_MODIFY_MAPPING.items():
+ key = re.sub(old_pattern, new_pattern, key)
+
+ new_state_dict[key] = value
+ return new_state_dict
+
+
+def convert_model(vq_model_id, llm_model_id, output_dir, hub_model_id=None, test_inference=False):
+ os.makedirs(output_dir, exist_ok=True)
+
+ # Convert and save processor
+ tokenizer_tiktoken = AutoTokenizer.from_pretrained(llm_model_id, trust_remote_code=True)
+ convert_tiktoken(tokenizer_tiktoken, output_dir)
+ extra_special_tokens = {
+ "image_token": "",
+ "boi_token": "<|image start|>",
+ "eoi_token": "<|image end|>",
+ "image_wrapper_token": "<|image token|>",
+ "eof_token": "<|extra_201|>",
+ }
+ tokenizer_converted = AutoTokenizer.from_pretrained(output_dir, extra_special_tokens=extra_special_tokens)
+ tokenizer_converted.padding_side = "left"
+
+ image_processor = Emu3ImageProcessor.from_pretrained(vq_model_id)
+ processor = Emu3Processor(image_processor, tokenizer_converted, chat_template=CHAT_TEMPLATE)
+ processor.save_pretrained(output_dir)
+
+ # load models
+ model_llm = AutoModelForCausalLM.from_pretrained(
+ llm_model_id,
+ trust_remote_code=True,
+ )
+ model_vqgan = AutoModel.from_pretrained(vq_model_id, trust_remote_code=True)
+ with open(f"{output_dir}/tokenizer.json", "r") as file:
+ tokenizer_config = json.load(file)
+ vocabulary_map = tokenizer_config["model"]["vocab"]
+
+ text_config = Emu3TextConfig(
+ max_position_embeddings=model_llm.config.max_position_embeddings,
+ rope_parameters={"rope_type": "default"},
+ )
+ config = Emu3Config(text_config=text_config, vocabulary_map=vocabulary_map)
+
+ with torch.device("meta"):
+ model = Emu3ForConditionalGeneration(config=config)
+ model.generation_config = GenerationConfig(
+ do_sample=True,
+ top_k=2048,
+ max_new_tokens=50_000,
+ pad_token_id=processor.tokenizer.pad_token_id,
+ eos_token_id=processor.tokenizer.eos_token_id,
+ )
+
+ state_dict = {}
+ state_dict = convert_state_dict_to_hf(model_llm.state_dict(), state_dict)
+ state_dict = convert_state_dict_to_hf(model_vqgan.state_dict(), state_dict)
+
+ model.load_state_dict(state_dict, assign=True, strict=True)
+ model.save_pretrained(output_dir)
+
+ if hub_model_id is not None:
+ model.push_to_hub(hub_model_id)
+ processor.push_to_hub(hub_model_id)
+
+ if test_inference and llm_model_id.endswith("Chat"):
+ # Short inference on a few examples to check if generation makes sense
+ print("Loading the checkpoint in a Emu3 model...")
+ print("*" * 100)
+ model = Emu3ForConditionalGeneration.from_pretrained(output_dir, dtype=torch.bfloat16, device_map="auto")
+ processor = Emu3Processor.from_pretrained(output_dir)
+
+ conversation = [
+ {
+ "role": "system",
+ "content": [
+ {"type": "text", "text": "You are a helpful assistant."},
+ ],
+ },
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "Please tell me about this art work and its artist."},
+ {"type": "image"},
+ ],
+ },
+ ]
+ prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+
+ url = "https://uploads4.wikiart.org/images/paul-klee/death-for-the-idea-1915.jpg!Large.jpg"
+ with httpx.stream("GET", url) as response:
+ image = Image.open(BytesIO(response.read()))
+ inputs = processor(images=image, text=prompt, return_tensors="pt").to(model.device, torch.bfloat16)
+ length = inputs.input_ids.shape[1]
+
+ out = model.generate(**inputs, max_new_tokens=40, do_sample=False)
+ generated_text = processor.batch_decode(out[:, length:], skip_special_tokens=True)[0]
+
+ print(f"Generation for single-image: {generated_text}")
+ print("*" * 100)
+ elif test_inference and llm_model_id.endswith("Gen"):
+ processor = Emu3Processor.from_pretrained(output_dir)
+ model = Emu3ForConditionalGeneration.from_pretrained(output_dir, dtype=torch.bfloat16, device_map="auto")
+
+ inputs = processor(
+ text=[
+ "a portrait of young girl. masterpiece, film grained, best quality.",
+ "a dog running under the rain",
+ ],
+ padding=True,
+ return_tensors="pt",
+ return_for_image_generation=True,
+ )
+ inputs = inputs.to(device="cuda:0", dtype=torch.bfloat16)
+
+ neg_prompt = "lowres, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry."
+ neg_inputs = processor(text=[neg_prompt] * 2, return_tensors="pt").to(device="cuda:0")
+
+ image_sizes = inputs.pop("image_sizes")
+ HEIGHT, WIDTH = image_sizes[0]
+ VISUAL_TOKENS = model.vocabulary_mapping.image_tokens
+
+ def prefix_allowed_tokens_fn(batch_id, input_ids):
+ height, width = HEIGHT, WIDTH
+ visual_tokens = VISUAL_TOKENS
+ image_token_id = processor.tokenizer.encode("<|image token|>", return_tensors="pt")[0].to(model.device)
+ eoi_token_id = processor.tokenizer.encode("<|image end|>", return_tensors="pt")[0]
+ eos_token_id = processor.tokenizer.encode("<|extra_204|>", return_tensors="pt")[0]
+ pad_token_id = processor.tokenizer.encode("<|endoftext|>", return_tensors="pt")[0]
+ eol_token_id = processor.tokenizer.encode("<|extra_200|>", return_tensors="pt")[0]
+ eof_token_id = processor.tokenizer.encode("<|extra_201|>", return_tensors="pt")[0]
+
+ position = torch.nonzero(input_ids == image_token_id, as_tuple=True)[0][0]
+ offset = input_ids.shape[0] - position
+ if offset % (width + 1) == 0:
+ return (eol_token_id,)
+ elif offset == (width + 1) * height + 1:
+ return (eof_token_id,)
+ elif offset == (width + 1) * height + 2:
+ return (eoi_token_id,)
+ elif offset == (width + 1) * height + 3:
+ return (eos_token_id,)
+ elif offset > (width + 1) * height + 3:
+ return (pad_token_id,)
+ else:
+ return visual_tokens
+
+ out = model.generate(
+ **inputs,
+ prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,
+ negative_prompt_ids=neg_inputs.input_ids,
+ negative_prompt_attention_mask=neg_inputs.attention_mask,
+ )
+
+ image = model.decode_image_tokens(out[:, inputs.input_ids.shape[1] :], height=HEIGHT, width=WIDTH)
+ images = processor.postprocess(
+ list(image.float()), return_tensors="PIL.Image.Image"
+ ) # internally we convert to np but it's not supported in bf16 precision
+ for i, image in enumerate(images["pixel_values"]):
+ image.save(f"result_{i}.png")
+
+
+def main():
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--vq_model_id",
+ help="Model ID of Emu3 VQ-VAE on the hub",
+ default="BAAI/Emu3-VisionTokenizer",
+ )
+ parser.add_argument(
+ "--llm_model_id",
+ help="Model ID of Emu3 bacbone LLM on the hub",
+ default="BAAI/Emu3-Chat",
+ )
+ parser.add_argument(
+ "--output_dir",
+ help="Location to write HF model",
+ )
+ parser.add_argument(
+ "--hub_model_id",
+ help="Model ID in the hub where to push the model.",
+ )
+ parser.add_argument(
+ "--test_inference",
+ action="store_true",
+ help="Whether to load the model for generation to test it's converted correctly.",
+ )
+ args = parser.parse_args()
+ convert_model(
+ vq_model_id=args.vq_model_id,
+ llm_model_id=args.llm_model_id,
+ output_dir=args.output_dir,
+ hub_model_id=args.hub_model_id,
+ test_inference=args.test_inference,
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/third_party/transformers/src/transformers/models/emu3/image_processing_emu3.py b/third_party/transformers/src/transformers/models/emu3/image_processing_emu3.py
new file mode 100644
index 0000000000000000000000000000000000000000..b159d457143b1f69814e6db804244698c651e065
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/emu3/image_processing_emu3.py
@@ -0,0 +1,534 @@
+# Copyright 2024 HuggingFace Inc. team. All rights reserved.
+#
+#
+# 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.
+
+import math
+from collections.abc import Iterable
+
+import numpy as np
+
+from ...image_processing_utils import BaseImageProcessor, BatchFeature
+from ...image_transforms import convert_to_rgb, pad, resize, to_channel_dimension_format
+from ...image_utils import (
+ OPENAI_CLIP_MEAN,
+ OPENAI_CLIP_STD,
+ ChannelDimension,
+ ImageInput,
+ PILImageResampling,
+ get_image_size,
+ infer_channel_dimension_format,
+ is_scaled_image,
+ make_flat_list_of_images,
+ make_nested_list_of_images,
+ to_numpy_array,
+ valid_images,
+ validate_preprocess_arguments,
+)
+from ...processing_utils import ImagesKwargs
+from ...utils import TensorType, is_vision_available, logging
+
+
+if is_vision_available():
+ from PIL import Image
+
+logger = logging.get_logger(__name__)
+
+
+class Emu3ImageProcessorKwargs(ImagesKwargs, total=False):
+ """
+ ratio (`str`, *optional*, defaults to `"1:1"`):
+ The ratio of the image to resize the image.
+ image_area (`int`, *optional*, defaults to `518400`):
+ The area of the image to resize the image.
+ """
+
+ ratio: str
+ image_area: int
+
+
+def smart_resize(
+ height: int, width: int, factor: int = 28, min_pixels: int = 56 * 56, max_pixels: int = 14 * 14 * 4 * 1280
+):
+ """Rescales the image so that the following conditions are met:
+
+ 1. Both dimensions (height and width) are divisible by 'factor'.
+
+ 2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+ 3. The aspect ratio of the image is maintained as closely as possible.
+
+ """
+ if max(height, width) / min(height, width) > 200:
+ raise ValueError(
+ f"absolute aspect ratio must be smaller than 200, got {max(height, width) / min(height, width)}"
+ )
+ h_bar = round(height / factor) * factor
+ w_bar = round(width / factor) * factor
+ if h_bar * w_bar > max_pixels:
+ beta = math.sqrt((height * width) / max_pixels)
+ h_bar = max(factor, math.floor(height / beta / factor) * factor)
+ w_bar = max(factor, math.floor(width / beta / factor) * factor)
+ elif h_bar * w_bar < min_pixels:
+ beta = math.sqrt(min_pixels / (height * width))
+ h_bar = math.ceil(height * beta / factor) * factor
+ w_bar = math.ceil(width * beta / factor) * factor
+ return h_bar, w_bar
+
+
+class Emu3ImageProcessor(BaseImageProcessor):
+ r"""
+ Constructs a Emu3 image processor that dynamically resizes images based on the original images.
+
+ Args:
+ do_resize (`bool`, *optional*, defaults to `True`):
+ Whether to resize the image's (height, width) dimensions.
+ resample (`PILImageResampling`, *optional*, defaults to `Resampling.BICUBIC`):
+ Resampling filter to use when resizing the image.
+ do_rescale (`bool`, *optional*, defaults to `True`):
+ Whether to rescale the image by the specified scale `rescale_factor`.
+ rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
+ Scale factor to use if rescaling the image.
+ do_normalize (`bool`, *optional*, defaults to `True`):
+ Whether to normalize the image.
+ image_mean (`float` or `list[float]`, *optional*, defaults to `[0.48145466, 0.4578275, 0.40821073]`):
+ Mean to use if normalizing the image. This is a float or list of floats for each channel in the image.
+ image_std (`float` or `list[float]`, *optional*, defaults to `[0.26862954, 0.26130258, 0.27577711]`):
+ Standard deviation to use if normalizing the image. This is a float or list of floats for each channel in the image.
+ do_convert_rgb (`bool`, *optional*, defaults to `True`):
+ Whether to convert the image to RGB.
+ do_pad (`bool`, *optional*, defaults to `True`):
+ Whether to pad the image. If `True`, will pad the patch dimension of the images in the batch to the largest
+ number of patches in the batch. Padding will be applied to the bottom and right with zeros.
+ min_pixels (`int`, *optional*, defaults to `512 * 512`):
+ The min pixels of the image to resize the image.
+ max_pixels (`int`, *optional*, defaults to `1024 * 1024`):
+ The max pixels of the image to resize the image.
+ spatial_factor (`int`, *optional*, defaults to 8):
+ The spatial downsample factor the image will be downsampled in feature extracting phase
+ """
+
+ model_input_names = ["pixel_values", "image_sizes"]
+ valid_kwargs = Emu3ImageProcessorKwargs
+
+ def __init__(
+ self,
+ do_resize: bool = True,
+ resample: PILImageResampling = PILImageResampling.BICUBIC,
+ do_rescale: bool = True,
+ rescale_factor: int | float = 1 / 255,
+ do_normalize: bool = True,
+ image_mean: float | list[float] | None = None,
+ image_std: float | list[float] | None = None,
+ do_convert_rgb: bool = True,
+ do_pad: bool = True,
+ min_pixels: int = 512 * 512,
+ max_pixels: int = 1024 * 1024,
+ spatial_factor: int = 8,
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.do_resize = do_resize
+ self.resample = resample
+ self.do_rescale = do_rescale
+ self.rescale_factor = rescale_factor
+ self.do_normalize = do_normalize
+ self.image_mean = image_mean if image_mean is not None else OPENAI_CLIP_MEAN
+ self.image_std = image_std if image_std is not None else OPENAI_CLIP_STD
+ self.min_pixels = min_pixels
+ self.max_pixels = max_pixels
+ self.spatial_factor = spatial_factor
+ self.size = {"min_pixels": min_pixels, "max_pixels": max_pixels}
+ self.do_convert_rgb = do_convert_rgb
+
+ def _preprocess(
+ self,
+ images: ImageInput,
+ do_resize: bool | None = None,
+ resample: PILImageResampling | None = None,
+ do_rescale: bool | None = None,
+ rescale_factor: float | None = None,
+ do_normalize: bool | None = None,
+ image_mean: float | list[float] | None = None,
+ image_std: float | list[float] | None = None,
+ do_convert_rgb: bool | None = None,
+ data_format: ChannelDimension | None = ChannelDimension.FIRST,
+ input_data_format: str | ChannelDimension | None = None,
+ ):
+ """
+ Preprocess an image or batch of images.
+
+ Args:
+ images (`ImageInput`):
+ Image or batch of images to preprocess. Expects pixel values ranging from 0 to 255. If pixel values range from 0 to 1, set `do_rescale=False`.
+ vision_info (`list[Dict]`, *optional*):
+ Optional list of dictionaries containing additional information about vision inputs.
+ do_resize (`bool`, *optional*, defaults to `self.do_resize`):
+ Whether to resize the image.
+ resample (`PILImageResampling`, *optional*, defaults to `self.resample`):
+ Resampling filter to use if resizing the image. This can be one of the `PILImageResampling` enums.
+ do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
+ Whether to rescale the image.
+ rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+ Scale factor to use if rescaling the image.
+ do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
+ Whether to normalize the image.
+ image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`):
+ Mean to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image.
+ image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`):
+ Standard deviation to use if normalizing the image. Can be a float or a list of floats corresponding to the number of channels in the image.
+ do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
+ Whether to convert the image to RGB.
+ data_format (`ChannelDimension`, *optional*, defaults to `ChannelDimension.FIRST`):
+ The channel dimension format for the output image. Can be one of:
+ - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+ - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+ - Unset: Use the channel dimension format of the input image.
+ input_data_format (`ChannelDimension` or `str`, *optional*):
+ The channel dimension format for the input image. Can be one of:
+ - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+ - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+ - `"none"` or `ChannelDimension.NONE`: image in (height, width) format. - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+ """
+ images = make_flat_list_of_images(images)
+
+ if do_convert_rgb:
+ images = [convert_to_rgb(image) for image in images]
+
+ # All transformations expect numpy arrays.
+ images = [to_numpy_array(image) for image in images]
+
+ if is_scaled_image(images[0]) and do_rescale:
+ logger.warning_once(
+ "It looks like you are trying to rescale already rescaled images. If the input"
+ " images have pixel values between 0 and 1, set `do_rescale=False` to avoid rescaling them again."
+ )
+ if input_data_format is None:
+ # We assume that all images have the same channel dimension format.
+ input_data_format = infer_channel_dimension_format(images[0])
+
+ processed_images = []
+ for image in images:
+ if do_resize:
+ height, width = get_image_size(image, channel_dim=input_data_format)
+ resized_height, resized_width = smart_resize(
+ height,
+ width,
+ factor=self.spatial_factor,
+ min_pixels=self.min_pixels,
+ max_pixels=self.max_pixels,
+ )
+ image = resize(
+ image, size=(resized_height, resized_width), resample=resample, input_data_format=input_data_format
+ )
+
+ if do_rescale:
+ image = self.rescale(image, scale=rescale_factor, input_data_format=input_data_format)
+
+ if do_normalize:
+ image = self.normalize(
+ image=image, mean=image_mean, std=image_std, input_data_format=input_data_format
+ )
+
+ image = to_channel_dimension_format(image, data_format, input_channel_dim=input_data_format)
+ processed_images.append(image)
+
+ return processed_images
+
+ def _pad_for_batching(
+ self,
+ pixel_values: list[np.ndarray],
+ image_sizes: list[list[int]],
+ data_format: str | ChannelDimension | None = None,
+ input_data_format: str | ChannelDimension | None = None,
+ ):
+ """
+ Pads images on the `num_of_patches` dimension with zeros to form a batch of same number of patches.
+
+ Args:
+ pixel_values (`list[np.ndarray]`):
+ An array of pixel values of each images of shape (`batch_size`, `num_patches`, `image_in_3D`)
+ image_sizes (`list[list[int]]`):
+ A list of sizes for each image in `pixel_values` in (height, width) format.
+ data_format (`str` or `ChannelDimension`, *optional*):
+ The channel dimension format for the output image. Can be one of:
+ - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+ - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+ If unset, will use same as the input image.
+ input_data_format (`str` or `ChannelDimension`, *optional*):
+ The channel dimension format for the input image. Can be one of:
+ - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+ - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+ If unset, will use the inferred format of the input image.
+
+ Returns:
+ list[`np.ndarray`]: The padded images.
+ """
+
+ max_shape = (
+ max(size[0] for size in image_sizes),
+ max(size[1] for size in image_sizes),
+ )
+ pixel_values = [
+ pad(
+ image,
+ padding=((0, max_shape[0] - size[0]), (0, max_shape[1] - size[1])),
+ data_format=data_format,
+ input_data_format=input_data_format,
+ )
+ for image, size in zip(pixel_values, image_sizes)
+ ]
+ return pixel_values
+
+ def preprocess(
+ self,
+ images: ImageInput,
+ do_resize: bool | None = None,
+ size: dict[str, int] | None = None,
+ resample: PILImageResampling | None = None,
+ do_rescale: bool | None = None,
+ rescale_factor: float | None = None,
+ do_normalize: bool | None = None,
+ image_mean: float | list[float] | None = None,
+ image_std: float | list[float] | None = None,
+ do_convert_rgb: bool | None = None,
+ do_pad: bool = True,
+ return_tensors: str | TensorType | None = None,
+ data_format: ChannelDimension | None = ChannelDimension.FIRST,
+ input_data_format: str | ChannelDimension | None = None,
+ ):
+ """
+ Args:
+ images (`ImageInput`):
+ Image to preprocess. Expects a single or batch of images with pixel values ranging from 0 to 255. If
+ passing in images with pixel values between 0 and 1, set `do_rescale=False`.
+ do_resize (`bool`, *optional*, defaults to `self.do_resize`):
+ Whether to resize the image.
+ size (`dict[str, int]`, *optional*, defaults to `self.size`):
+ Size of the image after resizing. Shortest edge of the image is resized to size["shortest_edge"], with
+ the longest edge resized to keep the input aspect ratio.
+ resample (`int`, *optional*, defaults to `self.resample`):
+ Resampling filter to use if resizing the image. This can be one of the enum `PILImageResampling`. Only
+ has an effect if `do_resize` is set to `True`.
+ do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
+ Whether to rescale the image.
+ rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+ Rescale factor to rescale the image by if `do_rescale` is set to `True`.
+ do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
+ Whether to normalize the image.
+ image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`):
+ Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
+ image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`):
+ Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to
+ `True`.
+ do_convert_rgb (`bool`, *optional*, defaults to `self.do_convert_rgb`):
+ Whether to convert the image to RGB.
+ do_pad (`bool`, *optional*, defaults to `True`):
+ Whether to pad the image. If `True`, will pad the patch dimension of the images in the batch to the largest
+ number of patches in the batch. Padding will be applied to the bottom and right with zeros.
+ return_tensors (`str` or `TensorType`, *optional*):
+ The type of tensors to return. Can be one of:
+ - Unset: Return a list of `np.ndarray`.
+ - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
+ - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
+ data_format (`ChannelDimension` or `str`, *optional*, defaults to `ChannelDimension.FIRST`):
+ The channel dimension format for the output image. Can be one of:
+ - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+ - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+ - Unset: Use the channel dimension format of the input image.
+ input_data_format (`ChannelDimension` or `str`, *optional*):
+ The channel dimension format for the input image. If unset, the channel dimension format is inferred
+ from the input image. Can be one of:
+ - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+ - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+ - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+
+ """
+ do_resize = do_resize if do_resize is not None else self.do_resize
+ size = size if size is not None else self.size
+ resample = resample if resample is not None else self.resample
+ do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+ rescale_factor = rescale_factor if rescale_factor is not None else self.rescale_factor
+ do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+ image_mean = image_mean if image_mean is not None else self.image_mean
+ image_std = image_std if image_std is not None else self.image_std
+ do_convert_rgb = do_convert_rgb if do_convert_rgb is not None else self.do_convert_rgb
+ do_pad = do_pad if do_pad is not None else self.do_pad
+
+ if images is not None:
+ images = self.fetch_images(images)
+ images = make_nested_list_of_images(images)
+
+ if images is not None and not valid_images(images):
+ raise ValueError("Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, or torch.Tensor")
+
+ validate_preprocess_arguments(
+ rescale_factor=rescale_factor,
+ do_normalize=do_normalize,
+ image_mean=image_mean,
+ image_std=image_std,
+ do_resize=do_resize,
+ size=size,
+ resample=resample,
+ )
+
+ pixel_values = []
+ for image in images:
+ if image:
+ image = self._preprocess(
+ image,
+ do_resize=do_resize,
+ resample=resample,
+ do_rescale=do_rescale,
+ rescale_factor=rescale_factor,
+ do_normalize=do_normalize,
+ image_mean=image_mean,
+ image_std=image_std,
+ data_format=data_format,
+ do_convert_rgb=do_convert_rgb,
+ input_data_format=input_data_format,
+ )
+ pixel_values.extend(image)
+
+ image_sizes = [get_image_size(image, input_data_format) for image in pixel_values]
+ if do_pad:
+ pixel_values = self._pad_for_batching(pixel_values, image_sizes)
+ pixel_values = np.array(pixel_values)
+
+ return BatchFeature(
+ data={"pixel_values": pixel_values, "image_sizes": image_sizes}, tensor_type=return_tensors
+ )
+
+ def postprocess(
+ self,
+ images: ImageInput,
+ do_rescale: bool | None = None,
+ rescale_factor: float | None = None,
+ do_normalize: bool | None = None,
+ image_mean: float | list[float] | None = None,
+ image_std: float | list[float] | None = None,
+ return_tensors: str | TensorType = "PIL.Image.Image",
+ input_data_format: str | ChannelDimension | None = None,
+ ):
+ """
+ Postprocess an image or batch of images tensor. Postprocess is the reverse process of preprocess.
+ The parameters should be same as in preprocess.
+ Args:
+ images (`ImageInput`):
+ Image to postprocess. Expects a single or batch of images with pixel values ranging from -1 to 1.
+ do_rescale (`bool`, *optional*, defaults to `self.do_rescale`):
+ Whether to rescale the image.
+ rescale_factor (`float`, *optional*, defaults to `self.rescale_factor`):
+ Rescale factor to rescale the image by if `do_rescale` is set to `True`.
+ do_normalize (`bool`, *optional*, defaults to `self.do_normalize`):
+ Whether to normalize the image.
+ image_mean (`float` or `list[float]`, *optional*, defaults to `self.image_mean`):
+ Image mean to use for normalization. Only has an effect if `do_normalize` is set to `True`.
+ image_std (`float` or `list[float]`, *optional*, defaults to `self.image_std`):
+ Image standard deviation to use for normalization. Only has an effect if `do_normalize` is set to `True`.
+ return_tensors (`str` or `TensorType`, *optional*):
+ The type of tensors to return. Can be one of:
+ - Unset: Return a list of `np.ndarray`.
+ - `TensorType.PYTORCH` or `'pt'`: Return a batch of type `torch.Tensor`.
+ - `TensorType.NUMPY` or `'np'`: Return a batch of type `np.ndarray`.
+ input_data_format (`ChannelDimension` or `str`, *optional*):
+ The channel dimension format for the input image. If unset, the channel dimension format is inferred
+ from the input image. Can be one of:
+ - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+ - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+ - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+ """
+ do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+ rescale_factor = 1.0 / self.rescale_factor if rescale_factor is None else rescale_factor
+ do_normalize = do_normalize if do_normalize is not None else self.do_normalize
+ image_mean = image_mean if image_mean is not None else self.image_mean
+ image_std = image_std if image_std is not None else self.image_std
+
+ images = make_flat_list_of_images(images)
+ if isinstance(images[0], Image.Image):
+ return images if len(images) > 1 else images[0]
+
+ if input_data_format is None:
+ # We assume that all images have the same channel dimension format.
+ input_data_format = infer_channel_dimension_format(images[0])
+
+ pixel_values = []
+ for image in images:
+ image = to_numpy_array(image)
+ if do_normalize:
+ image = self.unnormalize(
+ image=image, image_mean=image_mean, image_std=image_std, input_data_format=input_data_format
+ )
+
+ if do_rescale:
+ image = self.rescale(image, scale=rescale_factor, input_data_format=input_data_format)
+ image = image.clip(0, 255).astype(np.uint8)
+
+ if do_normalize and do_rescale and return_tensors == "PIL.Image.Image":
+ image = to_channel_dimension_format(image, ChannelDimension.LAST, input_channel_dim=input_data_format)
+ pixel_values.append(Image.fromarray(image))
+ else:
+ pixel_values.extend(image)
+
+ data = {"pixel_values": pixel_values}
+ return_tensors = return_tensors if return_tensors != "PIL.Image.Image" else None
+
+ return BatchFeature(data=data, tensor_type=return_tensors)
+
+ def unnormalize(
+ self,
+ image: np.ndarray,
+ image_mean: float | Iterable[float],
+ image_std: float | Iterable[float],
+ input_data_format: str | ChannelDimension | None = None,
+ ) -> np.ndarray:
+ """
+ Unnormalizes `image` using the mean and standard deviation specified by `mean` and `std`.
+ image = (image * image_std) + image_mean
+ Args:
+ image (`torch.Tensor` of shape `(batch_size, num_channels, image_size, image_size)` or `(num_channels, image_size, image_size)`):
+ Batch of pixel values to postprocess.
+ image_mean (`float` or `Iterable[float]`):
+ The mean to use for unnormalization.
+ image_std (`float` or `Iterable[float]`):
+ The standard deviation to use for unnormalization.
+ input_data_format (`ChannelDimension` or `str`, *optional*):
+ The channel dimension format for the input image. If unset, the channel dimension format is inferred
+ from the input image. Can be one of:
+ - `"channels_first"` or `ChannelDimension.FIRST`: image in (num_channels, height, width) format.
+ - `"channels_last"` or `ChannelDimension.LAST`: image in (height, width, num_channels) format.
+ - `"none"` or `ChannelDimension.NONE`: image in (height, width) format.
+ """
+ num_channels = 3
+
+ if isinstance(image_mean, Iterable):
+ if len(image_mean) != num_channels:
+ raise ValueError(f"mean must have {num_channels} elements if it is an iterable, got {len(image_mean)}")
+ else:
+ image_mean = [image_mean] * num_channels
+
+ if isinstance(image_std, Iterable):
+ if len(image_std) != num_channels:
+ raise ValueError(f"std must have {num_channels} elements if it is an iterable, got {len(image_std)}")
+ else:
+ image_std = [image_std] * num_channels
+
+ rev_image_mean = tuple(-mean / std for mean, std in zip(image_mean, image_std))
+ rev_image_std = tuple(1 / std for std in image_std)
+ image = self.normalize(
+ image=image, mean=rev_image_mean, std=rev_image_std, input_data_format=input_data_format
+ )
+ return image
+
+
+__all__ = ["Emu3ImageProcessor"]
diff --git a/third_party/transformers/src/transformers/models/emu3/modeling_emu3.py b/third_party/transformers/src/transformers/models/emu3/modeling_emu3.py
new file mode 100644
index 0000000000000000000000000000000000000000..20dd16f9ffb19caaacde8721a26f16553d3afda1
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/emu3/modeling_emu3.py
@@ -0,0 +1,1664 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/emu3/modular_emu3.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_emu3.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2024 HuggingFace Inc. team. All rights reserved.
+#
+#
+# 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.
+
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+from functools import cached_property
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
+from ...masking_utils import create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, CausalLMOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check
+from ...utils.generic import maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_emu3 import Emu3Config, Emu3TextConfig, Emu3VQVAEConfig
+
+
+@dataclass
+@auto_docstring
+class Emu3VQVAEModelOutput(BaseModelOutputWithPooling):
+ r"""
+ image_tokens (`torch.LongTensor` of shape `(batch_size, config.vocab_size`):
+ Indices of the image tokens predicted by the VQ-VAE model.
+ """
+
+ image_tokens: torch.LongTensor | None = None
+
+
+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)
+
+
+@use_kernel_func_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+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)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class Emu3Attention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: Emu3Config, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = True
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Emu3RMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ Emu3RMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class Emu3MLP(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=config.mlp_bias)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, x):
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+class Emu3DecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: Emu3Config, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = Emu3Attention(config=config, layer_idx=layer_idx)
+
+ self.mlp = Emu3MLP(config)
+ self.input_layernorm = Emu3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = Emu3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.dropout = nn.Dropout(config.attention_dropout)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + self.dropout(hidden_states)
+
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + self.dropout(hidden_states)
+ return hidden_states
+
+
+class Emu3VQVAEVectorQuantizer(nn.Module):
+ """
+ A module for vector quantization using learned embedding vectors.
+
+ This module implements the quantization process similar to te one described in
+ the VQ-VAE (Vector Quantized Variational AutoEncoder) paper. It quantizes continuous
+ input vectors into discrete codebook vectors, which are learned during training.
+ Current implementation improves over previous ones by avoiding costly matrix multiplications
+ and allowing for post-hoc remapping of indices.
+ """
+
+ def __init__(self, config: Emu3VQVAEConfig):
+ super().__init__()
+ self.embedding = nn.Embedding(config.codebook_size, config.embed_dim)
+ self.embedding.weight.data.uniform_(-1.0 / config.codebook_size, 1.0 / config.codebook_size)
+
+ def forward(self, hidden_state: torch.Tensor):
+ batch_size, temporal, channels, height, width = hidden_state.shape
+ hidden_state = hidden_state.permute(0, 1, 3, 4, 2).contiguous()
+ hidden_state_flattened = hidden_state.view(-1, channels)
+
+ # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+ hidden_state_sum = torch.sum(hidden_state_flattened**2, dim=1, keepdim=True)
+ embedding_sum = torch.sum(self.embedding.weight**2, dim=1)
+
+ # "bd,dn->bn",
+ distances = 2 * torch.matmul(hidden_state_flattened, self.embedding.weight.transpose(0, 1))
+ distances = hidden_state_sum + embedding_sum - distances
+
+ min_encoding_indices = torch.argmin(distances, dim=1)
+ min_encoding_indices = min_encoding_indices.view(batch_size, temporal, height, width)
+ return min_encoding_indices
+
+
+class Emu3VQVAEEncoderConvDownsample(nn.Module):
+ def __init__(self, in_channels):
+ super().__init__()
+ self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)
+
+ def forward(self, hidden_states):
+ # no asymmetric padding in torch conv, must do it ourselves
+ hidden_states = F.pad(hidden_states, pad=(0, 1, 0, 1), mode="constant", value=0)
+ hidden_states = self.conv(hidden_states)
+ return hidden_states
+
+
+class Emu3VQVAEEncoderConvUpsample(nn.Module):
+ def __init__(self, in_channels):
+ super().__init__()
+ self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+ def forward(self, hidden_states):
+ hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
+ hidden_states = self.conv(hidden_states)
+ return hidden_states
+
+
+class Emu3VQVAEConv3d(nn.Module):
+ def __init__(
+ self,
+ in_channel: int,
+ out_channel: int,
+ kernel_size: tuple[int],
+ stride: tuple[int],
+ ):
+ super().__init__()
+
+ padding_sizes = [one_kernel - one_stride for one_kernel, one_stride in zip(kernel_size[1:], stride[1:])]
+ self.padding = ()
+ for pad_size in padding_sizes[::-1]:
+ self.padding += (pad_size // 2 + pad_size % 2, pad_size // 2)
+ self.padding += (2, 0)
+
+ self.conv = nn.Conv3d(
+ in_channel,
+ out_channel,
+ kernel_size,
+ stride=stride,
+ )
+
+ def forward(self, hidden_states: torch.Tensor):
+ hidden_states = F.pad(hidden_states, self.padding)
+ hidden_states = self.conv(hidden_states)
+ return hidden_states
+
+
+class Emu3VQVAESpatialNorm(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ ):
+ super().__init__()
+ self.norm_layer = nn.GroupNorm(
+ num_channels=out_channels,
+ num_groups=32,
+ eps=1e-6,
+ affine=True,
+ )
+
+ self.conv_y = nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ )
+ self.conv_b = nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ )
+
+ def forward(self, hidden_states: torch.Tensor, quant_states: torch.Tensor):
+ quant_states = F.interpolate(quant_states, size=hidden_states.shape[-2:], mode="nearest")
+ hidden_states = self.norm_layer(hidden_states)
+ hidden_states = hidden_states * self.conv_y(quant_states) + self.conv_b(quant_states)
+ return hidden_states
+
+
+class Emu3VQVAETemporalUpsample(nn.Module):
+ def __init__(
+ self,
+ in_channel: int,
+ out_channel: int,
+ ):
+ super().__init__()
+ self.conv = Emu3VQVAEConv3d(
+ in_channel,
+ out_channel,
+ kernel_size=(3, 3, 3),
+ stride=(1, 1, 1),
+ )
+
+ def forward(self, hidden_states: torch.Tensor):
+ batch_size, channels, temporal, height, width = hidden_states.shape
+ hidden_states = hidden_states.permute(0, 1, 3, 4, 2).contiguous().view(batch_size, -1, temporal)
+ hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
+ hidden_states = hidden_states.view(batch_size, channels, height, width, -1).permute(0, 1, 4, 2, 3).contiguous()
+ hidden_states = self.conv(hidden_states)
+ return hidden_states
+
+
+class Emu3VQVAETemporalDownsample(nn.Module):
+ def __init__(
+ self,
+ in_channel: int,
+ out_channel: int,
+ ):
+ super().__init__()
+ self.conv = Emu3VQVAEConv3d(
+ in_channel,
+ out_channel,
+ kernel_size=(4, 3, 3),
+ stride=(2, 1, 1),
+ )
+
+ def forward(self, hidden_states: torch.Tensor):
+ hidden_states = self.conv(hidden_states)
+ return hidden_states
+
+
+class Emu3VQVAETemporalResnetBlock(nn.Module):
+ def __init__(
+ self,
+ in_channels,
+ out_channels=None,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ self.out_channels = in_channels if out_channels is None else out_channels
+
+ self.norm1 = nn.BatchNorm3d(in_channels)
+ self.conv1 = Emu3VQVAEConv3d(
+ in_channels,
+ out_channels,
+ kernel_size=(3, 3, 3),
+ stride=(1, 1, 1),
+ )
+ self.norm2 = nn.BatchNorm3d(out_channels)
+ self.conv2 = Emu3VQVAEConv3d(
+ out_channels,
+ out_channels,
+ kernel_size=(3, 3, 3),
+ stride=(1, 1, 1),
+ )
+ if self.in_channels != self.out_channels:
+ self.nin_shortcut = nn.Conv3d(
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ )
+
+ def forward(self, hidden_states):
+ residual = hidden_states
+ hidden_states = self.norm1(hidden_states)
+ hidden_states *= torch.sigmoid(hidden_states)
+ hidden_states = self.conv1(hidden_states)
+
+ hidden_states = self.norm2(hidden_states)
+ hidden_states *= torch.sigmoid(hidden_states)
+ hidden_states = self.conv2(hidden_states)
+
+ if self.in_channels != self.out_channels:
+ residual = self.nin_shortcut(residual)
+
+ return residual + hidden_states
+
+
+class Emu3VQVAEResnetBlock(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int | None = None,
+ quant_channels: int | None = None,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ out_channels = in_channels if out_channels is None else out_channels
+ self.out_channels = out_channels
+ self.quant_channels = quant_channels
+
+ if quant_channels is None:
+ self.norm1 = nn.GroupNorm(num_channels=in_channels, num_groups=32, eps=1e-6, affine=True)
+ self.norm2 = nn.GroupNorm(num_channels=out_channels, num_groups=32, eps=1e-6, affine=True)
+ else:
+ self.norm1 = Emu3VQVAESpatialNorm(quant_channels, in_channels)
+ self.norm2 = Emu3VQVAESpatialNorm(quant_channels, out_channels)
+
+ self.conv1 = nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ self.conv2 = nn.Conv2d(
+ out_channels,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ if self.in_channels != self.out_channels:
+ self.nin_shortcut = nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ )
+
+ def forward(self, hidden_states: torch.Tensor, quant_channels: torch.Tensor | None = None):
+ norm_args = () if self.quant_channels is None else (quant_channels,)
+
+ residual = hidden_states
+ hidden_states = self.norm1(hidden_states, *norm_args)
+ hidden_states *= torch.sigmoid(hidden_states)
+ hidden_states = self.conv1(hidden_states)
+
+ hidden_states = self.norm2(hidden_states, *norm_args)
+ hidden_states *= torch.sigmoid(hidden_states)
+ hidden_states = self.conv2(hidden_states)
+
+ if self.in_channels != self.out_channels:
+ residual = self.nin_shortcut(residual)
+
+ return residual + hidden_states
+
+
+class Emu3VQVAEAttentionBlock(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: Emu3VQVAEConfig):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = self.embed_dim // self.num_heads
+ if self.head_dim * self.num_heads != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+ f" {self.num_heads})."
+ )
+ self.scale = self.head_dim**-0.5
+ self.dropout = config.attention_dropout
+ self.is_causal = False
+
+ self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+ # for compatibility with the attention interface
+ self.num_key_value_groups = 1
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Input shape: Batch x Time x Channel"""
+
+ batch_size, seq_length, embed_dim = hidden_states.shape
+
+ queries = self.q_proj(hidden_states)
+ keys = self.k_proj(hidden_states)
+ values = self.v_proj(hidden_states)
+
+ queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+ keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+ values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ queries,
+ keys,
+ values,
+ attention_mask,
+ is_causal=self.is_causal,
+ scaling=self.scale,
+ dropout=0.0 if not self.training else self.dropout,
+ )
+
+ attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+class Emu3VQVAEGroupNorm(nn.GroupNorm):
+ """
+ Same as the torch GroupNorm with the only difference that this ones accepts
+ an optional kwarg `quant_states` which is not used. This class makes it easier to
+ use SpatialNorm or GroupNorm without conditionals
+ """
+
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def forward(self, input, quant_states=None):
+ return F.group_norm(input, self.num_groups, self.weight, self.bias, self.eps)
+
+
+class Emu3VQVAEMiddleBlock(nn.Module):
+ def __init__(self, config, in_channels, quant_channels=None):
+ super().__init__()
+
+ self.block_1 = Emu3VQVAEResnetBlock(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ quant_channels=quant_channels,
+ )
+ self.attn_1 = Emu3VQVAEAttentionBlock(config)
+ if quant_channels is None:
+ self.attn_norm = Emu3VQVAEGroupNorm(num_channels=in_channels, num_groups=32, eps=1e-6, affine=True)
+ else:
+ self.attn_norm = Emu3VQVAESpatialNorm(quant_channels, in_channels)
+
+ self.block_2 = Emu3VQVAEResnetBlock(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ quant_channels=quant_channels,
+ )
+
+ def forward(self, hidden_states: torch.FloatTensor, quant_states: torch.FloatTensor | None = None):
+ hidden_states = self.block_1(hidden_states, quant_states)
+ residual = hidden_states
+ hidden_states = self.attn_norm(hidden_states, quant_states)
+ batch_size, channels, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)
+ hidden_states = self.attn_1(hidden_states)[0]
+ hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
+ hidden_states = residual + hidden_states
+ hidden_states = self.block_2(hidden_states, quant_states)
+ return hidden_states
+
+
+class Emu3VQVAEDownBlock(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ self.num_resolutions = len(config.channel_multiplier)
+ self.num_res_blocks = config.num_res_blocks
+ base_channels = config.base_channels
+ channel_multiplier = config.channel_multiplier
+
+ in_channel_multiplier = (1,) + tuple(channel_multiplier)
+ self.in_channel_multiplier = in_channel_multiplier
+ self.down = nn.ModuleList()
+ for i_level in range(self.num_resolutions):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ attn_norms = nn.ModuleList()
+ block_in = base_channels * in_channel_multiplier[i_level]
+ block_out = base_channels * channel_multiplier[i_level]
+ for i_block in range(self.num_res_blocks):
+ block.append(
+ Emu3VQVAEResnetBlock(
+ in_channels=block_in,
+ out_channels=block_out,
+ )
+ )
+ block_in = block_out
+ if config.attn_resolutions is not None and i_level in config.attn_resolutions:
+ attn.append(Emu3VQVAEAttentionBlock(config))
+ attn_norms.append(nn.GroupNorm(num_channels=block_in, num_groups=32, eps=1e-6, affine=True))
+
+ down = nn.Module()
+ down.block = block
+ down.attn = attn
+ down.attn_norms = attn_norms
+ if i_level != self.num_resolutions - 1:
+ down.downsample = Emu3VQVAEEncoderConvDownsample(block_in)
+ self.down.append(down)
+
+ def forward(self, hidden_states: torch.FloatTensor):
+ for i_level, blocks in enumerate(self.down):
+ for i_block in range(self.num_res_blocks):
+ hidden_states = blocks.block[i_block](hidden_states)
+ if len(blocks.attn) > 0:
+ residual = hidden_states
+ hidden_states = blocks.attn_norms[i_block](hidden_states)
+
+ batch_size, channels, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)
+ hidden_states = blocks.attn[i_block](hidden_states)[0]
+
+ hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
+ hidden_states = residual + hidden_states
+
+ if i_level != self.num_resolutions - 1:
+ hidden_states = blocks.downsample(hidden_states)
+
+ return hidden_states
+
+
+class Emu3VQVAEUpBlock(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ self.num_resolutions = len(config.channel_multiplier)
+ self.num_res_blocks = config.num_res_blocks
+
+ quant_channels = config.embed_dim
+ block_in = config.base_channels * config.channel_multiplier[-1]
+
+ self.up = nn.ModuleList()
+ for i_level in reversed(range(self.num_resolutions)):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ attn_norms = nn.ModuleList()
+ block_out = config.base_channels * config.channel_multiplier[i_level]
+ for i_block in range(self.num_res_blocks + 1):
+ block.append(
+ Emu3VQVAEResnetBlock(
+ in_channels=block_in,
+ out_channels=block_out,
+ quant_channels=quant_channels,
+ )
+ )
+ block_in = block_out
+ if i_level in config.attn_resolutions:
+ attn.append(Emu3VQVAEAttentionBlock(config))
+ attn_norms.append(Emu3VQVAESpatialNorm(quant_channels, block_in))
+
+ up = nn.Module()
+ up.block = block
+ up.attn = attn
+ up.attn_norms = attn_norms
+ if i_level != 0:
+ up.upsample = Emu3VQVAEEncoderConvUpsample(block_in)
+
+ self.up.insert(0, up)
+
+ def forward(self, hidden_states: torch.FloatTensor, quant_states: torch.FloatTensor):
+ for i_level, blocks in enumerate(self.up[::-1]):
+ for i_block in range(self.num_res_blocks + 1):
+ hidden_states = blocks.block[i_block](hidden_states, quant_states)
+ if len(blocks.attn) > 0:
+ residual = hidden_states
+ hidden_states = blocks.attn_norms[i_block](hidden_states, quant_states)
+
+ batch_size, channels, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)
+ hidden_states = blocks.attn[i_block](hidden_states)[0]
+
+ hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
+ hidden_states = residual + hidden_states
+ if i_level != len(self.up) - 1:
+ hidden_states = blocks.upsample(hidden_states)
+
+ return hidden_states
+
+
+class Emu3VQVAEEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ base_channels = config.base_channels
+ in_channels = config.in_channels
+ double_latent = config.double_latent
+ latent_channels = config.latent_channels
+ channel_multiplier = config.channel_multiplier
+ out_channels = 2 * latent_channels if double_latent else latent_channels
+ block_in = base_channels * channel_multiplier[-1]
+
+ self.conv_in = torch.nn.Conv2d(in_channels, base_channels, kernel_size=3, stride=1, padding=1)
+ self.down_block = Emu3VQVAEDownBlock(config)
+ self.middle_block = Emu3VQVAEMiddleBlock(config, block_in)
+
+ self.norm_out = torch.nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+ self.conv_out = torch.nn.Conv2d(
+ block_in,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ temporal_down_blocks = int(math.log2(config.temporal_downsample_factor))
+ self.time_conv = nn.ModuleList()
+ self.time_res_stack = nn.ModuleList()
+
+ for i in range(temporal_down_blocks):
+ conv = Emu3VQVAETemporalDownsample(out_channels, out_channels)
+ self.time_conv.append(conv)
+
+ for _ in range(config.num_res_blocks):
+ time_res_conv = Emu3VQVAETemporalResnetBlock(
+ in_channels=out_channels,
+ out_channels=out_channels,
+ )
+ self.time_res_stack.append(time_res_conv)
+
+ def forward(self, pixel_values: torch.LongTensor):
+ temporal_dim = pixel_values.shape[1]
+ pixel_values = pixel_values.reshape(-1, *pixel_values.shape[2:])
+
+ # downsampling & middle
+ hidden_states = self.conv_in(pixel_values)
+ hidden_states = self.down_block(hidden_states)
+ hidden_states = self.middle_block(hidden_states)
+
+ # end
+ hidden_states = self.norm_out(hidden_states)
+ hidden_states *= torch.sigmoid(hidden_states)
+ hidden_states = self.conv_out(hidden_states)
+
+ hidden_states = hidden_states.reshape(-1, temporal_dim, *hidden_states.shape[1:])
+ hidden_states = hidden_states.permute(0, 2, 1, 3, 4)
+
+ # temporal convs
+ for conv in self.time_conv:
+ hidden_states = conv(hidden_states)
+ hidden_states *= torch.sigmoid(hidden_states)
+
+ for layer in self.time_res_stack:
+ hidden_states = layer(hidden_states)
+
+ hidden_states = hidden_states.permute(0, 2, 1, 3, 4)
+
+ return hidden_states
+
+
+class Emu3VQVAEDecoder(nn.Module):
+ def __init__(self, config: Emu3VQVAEConfig):
+ super().__init__()
+
+ quant_channels = config.embed_dim
+ block_in = config.base_channels * config.channel_multiplier[-1]
+ self.time_res_stack = nn.ModuleList()
+ for _ in range(config.num_res_blocks):
+ time_res_conv = Emu3VQVAETemporalResnetBlock(
+ in_channels=config.latent_channels, out_channels=config.latent_channels
+ )
+ self.time_res_stack.append(time_res_conv)
+
+ temp_upsample_block_num = int(math.log2(config.temporal_downsample_factor))
+ self.time_conv = nn.ModuleList()
+ for i in range(temp_upsample_block_num):
+ conv = Emu3VQVAETemporalUpsample(config.latent_channels, config.latent_channels)
+ self.time_conv.append(conv)
+
+ self.conv_in = nn.Conv2d(
+ config.latent_channels,
+ block_in,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ self.middle_block = Emu3VQVAEMiddleBlock(config, block_in, quant_channels=quant_channels)
+ self.up_block = Emu3VQVAEUpBlock(config)
+
+ block_in = config.base_channels * config.channel_multiplier[0]
+ self.norm_out = Emu3VQVAESpatialNorm(quant_channels, block_in)
+ self.conv_out = nn.Conv2d(
+ block_in,
+ config.out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ def forward(self, hidden_states: torch.Tensor, quant_states: torch.Tensor):
+ hidden_quant_states = torch.cat((hidden_states, quant_states), dim=0)
+ hidden_quant_states = hidden_quant_states.permute(0, 2, 1, 3, 4)
+
+ # temporal convs
+ for layer in self.time_res_stack:
+ hidden_quant_states = layer(hidden_quant_states)
+
+ for layer in self.time_conv:
+ hidden_quant_states = layer(hidden_quant_states)
+ hidden_quant_states *= torch.sigmoid(hidden_quant_states)
+
+ hidden_quant_states = hidden_quant_states.permute(0, 2, 1, 3, 4)
+ hidden_states, quant_states = torch.chunk(hidden_quant_states, 2, dim=0)
+ hidden_states = hidden_states.reshape(-1, *hidden_states.shape[2:])
+ quant_states = quant_states.reshape(-1, *quant_states.shape[2:])
+
+ hidden_states = self.conv_in(hidden_states)
+
+ # middle & upsampling
+ hidden_states = self.middle_block(hidden_states, quant_states)
+ hidden_states = self.up_block(hidden_states, quant_states)
+
+ hidden_states = self.norm_out(hidden_states, quant_states)
+ hidden_states *= torch.sigmoid(hidden_states)
+ hidden_states = self.conv_out(hidden_states)
+
+ return hidden_states
+
+
+@auto_docstring(
+ custom_intro="""
+ The VQ-VAE model used in Emu3 for encoding/decoding images into discrete tokens.
+ This model follows the "Make-a-scene: Scene-based text-to-image generation with human priors" paper from
+ [ Oran Gafni, Adam Polyak, Oron Ashual, Shelly Sheynin, Devi Parikh, and Yaniv
+ Taigman](https://huggingface.co/papers/2203.13131).
+ """
+)
+class Emu3VQVAE(PreTrainedModel):
+ config: Emu3VQVAEConfig
+ base_model_prefix = "emuvideovq"
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+ _supports_sdpa = True
+ _supports_flash_attn = True
+ _supports_flex_attn = True
+ _supports_attention_backend = True
+ _no_split_modules = [
+ "Emu3VQVAETemporalResnetBlock",
+ "Emu3VQVAEAttentionBlock",
+ "Emu3VQVAEResnetBlock",
+ "Emu3VQVAEVectorQuantizer",
+ ]
+ _can_record_outputs = {
+ "hidden_states": [Emu3VQVAEResnetBlock, Emu3VQVAETemporalResnetBlock],
+ "attentions": Emu3VQVAEAttentionBlock,
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ if isinstance(module, (nn.Conv2d, nn.Conv3d)):
+ init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
+ if module.bias is not None:
+ fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(module.weight)
+ bound = 1 / math.sqrt(fan_in)
+ init.uniform_(module.bias, -bound, bound)
+ elif isinstance(module, nn.Linear):
+ init.kaiming_uniform_(module.weight, a=math.sqrt(5))
+ if module.bias is not None:
+ fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(module.weight)
+ bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+ init.uniform_(module.bias, -bound, bound)
+ elif isinstance(module, (nn.BatchNorm2d, nn.BatchNorm3d, nn.GroupNorm)):
+ init.constant_(module.weight, 1.0)
+ init.constant_(module.bias, 0.0)
+ if getattr(module, "running_mean", None) is not None:
+ init.zeros_(module.running_mean)
+ init.ones_(module.running_var)
+ init.zeros_(module.num_batches_tracked)
+ elif isinstance(module, nn.Embedding):
+ init.normal_(module.weight)
+ # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag
+ if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False):
+ init.zeros_(module.weight[module.padding_idx])
+
+ def __init__(self, config: Emu3VQVAEConfig):
+ super().__init__(config)
+
+ self.config = config
+
+ self.encoder = Emu3VQVAEEncoder(config)
+ self.decoder = Emu3VQVAEDecoder(config)
+ self.quantize = Emu3VQVAEVectorQuantizer(config)
+ self.vision_spatial_factor = 2 ** (len(config.channel_multiplier) - 1)
+
+ self.quant_conv = Emu3VQVAEConv3d(
+ config.latent_channels, config.embed_dim, kernel_size=(3, 1, 1), stride=(1, 1, 1)
+ )
+ self.post_quant_conv = Emu3VQVAEConv3d(
+ config.embed_dim, config.latent_channels, kernel_size=(3, 1, 1), stride=(1, 1, 1)
+ )
+ self.spatial_scale_factor = 2 ** (len(config.channel_multiplier) - 1)
+ self.eval() # Emu3's VQ model is frozen
+
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def encode(
+ self, pixel_values: torch.Tensor, image_sizes: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
+ ) -> Emu3VQVAEModelOutput:
+ is_image = pixel_values.ndim == 4
+ if is_image:
+ temporal = self.config.temporal_downsample_factor
+ batch_size, channels, height, width = pixel_values.shape
+ pixel_values = pixel_values.unsqueeze(1).repeat(1, temporal, 1, 1, 1)
+ else:
+ batch_size, temporal, channels, height, width = pixel_values.shape
+
+ hidden_states = self.encoder(pixel_values)
+
+ # b t c h w -> b c t h w
+ conv_hidden_states = hidden_states.permute(0, 2, 1, 3, 4)
+ conv_hidden_states = self.quant_conv(conv_hidden_states)
+
+ # b c t h w -> b t c h w
+ conv_hidden_states = conv_hidden_states.permute(0, 2, 1, 3, 4)
+ codes = self.quantize(conv_hidden_states)
+
+ image_tokens = codes.squeeze(1) if is_image else codes
+
+ image_tokens = [
+ single_image[: int(size[0] / self.vision_spatial_factor), : int(size[1] / self.vision_spatial_factor)]
+ for single_image, size in zip(image_tokens, image_sizes)
+ ]
+
+ return Emu3VQVAEModelOutput(
+ last_hidden_state=hidden_states,
+ image_tokens=image_tokens,
+ )
+
+ def decode(self, hidden_states: torch.Tensor):
+ is_image = hidden_states.ndim == 3
+ if is_image:
+ hidden_states = hidden_states.unsqueeze(1)
+
+ batch_size, temporal, height, width = hidden_states.shape
+ quant = self.quantize.embedding(hidden_states.flatten())
+
+ channels = quant.shape[-1]
+ quant = quant.view(batch_size, temporal, height, width, channels).permute(0, 4, 1, 2, 3).contiguous()
+ post_quant = self.post_quant_conv(quant)
+
+ quant = quant.permute(0, 2, 1, 3, 4)
+ post_quant = post_quant.permute(0, 2, 1, 3, 4)
+
+ video = self.decoder(post_quant, quant)
+ video = video.reshape(
+ batch_size,
+ temporal * self.config.temporal_downsample_factor,
+ self.config.out_channels,
+ height * self.spatial_scale_factor,
+ width * self.spatial_scale_factor,
+ )
+ return video[:, 0] if is_image else video
+
+
+class Emu3ImageVocabularyMapping:
+ """
+ A class for mapping discrete image tokens from VQGAN to BPE tokens.
+ """
+
+ def __init__(self, vocab_map):
+ self.vocab_map = vocab_map
+ self.eol_token_id = vocab_map.get("<|extra_200|>")
+ self.image_token_id = vocab_map.get("")
+
+ @cached_property
+ def image_tokens(self):
+ return sorted([val for name, val in self.vocab_map.items() if name.startswith("<|visual token")])
+
+ @cached_property
+ def image_tokens_str(self):
+ return sorted([name for name, val in self.vocab_map.items() if name.startswith("<|visual token")])
+
+ @cached_property
+ def img2bpe(self):
+ return {int(token[-8:-2]): self.vocab_map[token] for token in self.image_tokens_str}
+
+ @cached_property
+ def bpe2img(self):
+ return {v: k for k, v in self.img2bpe.items()}
+
+ @cached_property
+ def bpe2img_mapping_tensor(self):
+ mapping = torch.zeros(max(self.bpe2img.keys()) + 1, dtype=torch.int)
+ for k, v in self.bpe2img.items():
+ mapping[k] = v
+ return mapping
+
+ @cached_property
+ def img2bpe_mapping_tensor(self):
+ mapping = torch.zeros(max(self.img2bpe.keys()) + 1, dtype=torch.int)
+ for k, v in self.img2bpe.items():
+ mapping[k] = v
+ return mapping
+
+ def convert_img2bpe(self, img_batch: list[torch.Tensor]) -> torch.Tensor:
+ device = img_batch.device
+ eol_row = torch.ones((img_batch.shape[0], 1), dtype=torch.int) * self.eol_token_id
+ img_tokens = self.img2bpe_mapping_tensor[img_batch.to("cpu")]
+ img_tokens = torch.cat([img_tokens, eol_row], dim=-1)
+ return img_tokens.to(device)
+
+ def convert_bpe2img(self, img_batch: torch.Tensor) -> torch.Tensor:
+ device = img_batch.device
+ img_batch = img_batch[..., :-1] # remove last row of EOL tokens
+ img_tokens = self.bpe2img_mapping_tensor[img_batch.to("cpu")]
+ return img_tokens.to(device)
+
+
+@auto_docstring
+class Emu3PreTrainedModel(PreTrainedModel):
+ config: Emu3Config
+ base_model_prefix = "model"
+ input_modalities = ("image", "text")
+ supports_gradient_checkpointing = True
+ _no_split_modules = [
+ "Emu3DecoderLayer",
+ ]
+ _skip_keys_device_placement = ["past_key_values", "causal_mask"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+
+ _can_compile_fullgraph = True
+ _supports_flex_attn = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": Emu3DecoderLayer,
+ "attentions": Emu3Attention,
+ }
+
+
+class Emu3RotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: Emu3Config, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: Emu3Config | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@auto_docstring
+class Emu3TextModel(Emu3PreTrainedModel):
+ config: Emu3TextConfig
+
+ def __init__(self, config: Emu3TextConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList(
+ [Emu3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = Emu3RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = Emu3RotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_embeddings=position_embeddings,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class Emu3ForCausalLM(Emu3PreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+ _tp_plan = {"lm_head": "colwise_gather_output"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+ config: Emu3TextConfig
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = Emu3TextModel(config)
+ self.vocab_size = config.vocab_size
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> CausalLMOutputWithPast:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import Emu3Processor, Emu3ForConditionalGeneration
+ >>> import torch
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from PIL import Image
+
+ >>> model = Emu3ForCausalLM.from_pretrained("BAAI/Emu3-Chat-hf", dtype=torch.bfloat16)
+ >>> processor = Emu3Processor.from_pretrained("BAAI/Emu3-Chat-hf")
+
+ >>> inputs = processor(text=["Can you write me a poem about winter."], return_tensors="pt").to(model.device)
+
+ >>> generated_ids = model.generate(**inputs, max_new_tokens=100, do_sample=False)
+ >>> processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+ ```"""
+ outputs: BaseModelOutputWithPast = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+class Emu3Model(Emu3PreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.text_model = Emu3TextModel._from_config(config.text_config)
+ self.vqmodel = Emu3VQVAE(config.vq_config)
+ self.vocabulary_mapping = Emu3ImageVocabularyMapping(config.vocabulary_map)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.text_model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.text_model.set_input_embeddings(value)
+
+ def get_image_tokens(self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor) -> torch.LongTensor:
+ """
+ Tokenizes images into discrete tokens with VQGAN module. Converts
+ obtained image tokens into BPE tokens and wraps with "boi" and "eoi"
+ special tokens.
+
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_sizes (`torch.LongTensor` of shape `(batch_size, 2)`):
+ The sizes of the images in the batch, being (height, width) for each image.
+ """
+ vqmodel_outputs: Emu3VQVAEModelOutput = self.vqmodel.encode(pixel_values, image_sizes, return_dict=True)
+ bpe_tokens_list = [
+ self.vocabulary_mapping.convert_img2bpe(tokens).flatten() for tokens in vqmodel_outputs.image_tokens
+ ]
+ bpe_tokens = torch.cat(bpe_tokens_list)
+ return bpe_tokens
+
+ @can_return_tuple
+ @auto_docstring(
+ custom_intro="Tokenizes images into discrete tokens with VQGAN module and embeds them with text embeddings layer"
+ )
+ def get_image_features(
+ self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor, **kwargs: Unpack[TransformersKwargs]
+ ) -> tuple | Emu3VQVAEModelOutput:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)):
+ The tensors corresponding to the input images.
+ """
+ vqmodel_outputs: Emu3VQVAEModelOutput = self.vqmodel.encode(
+ pixel_values, image_sizes, return_dict=True, **kwargs
+ )
+ split_sizes = [
+ (height // self.vqmodel.vision_spatial_factor) * (width // self.vqmodel.vision_spatial_factor + 1)
+ for height, width in image_sizes
+ ]
+ bpe_tokens_list = [
+ self.vocabulary_mapping.convert_img2bpe(tokens).flatten() for tokens in vqmodel_outputs.image_tokens
+ ]
+ bpe_tokens = torch.cat(bpe_tokens_list)
+ image_embeddings = self.get_input_embeddings()(bpe_tokens)
+ image_features = torch.split(image_embeddings, split_sizes)
+ vqmodel_outputs.pooler_output = image_features
+
+ return vqmodel_outputs
+
+ @torch.no_grad()
+ def decode_image_tokens(self, image_tokens: torch.LongTensor, height: int, width: int):
+ """
+ Decodes generated image tokens from language model to continuous pixel values
+ with VQGAN module via upsampling.
+
+ Args:
+ image_tokens (`torch.LongTensor` of shape `(batch_size, num_of_tokens)`):
+ The tensors corresponding to the input images.
+ height (`int`):
+ Height of the generated image before upsampling.
+ width (`int`):
+ Width of the generated image before upsampling.
+ """
+ sequences = image_tokens[:, :-3].view(-1, height, width + 1)
+ image_tokens = self.vocabulary_mapping.convert_bpe2img(sequences)
+ image = self.vqmodel.decode(image_tokens)
+ return image
+
+ def get_placeholder_mask(
+ self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor
+ ):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+ equal to the length of multimodal features. If the lengths are different, an error is raised.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.vocabulary_mapping.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ else:
+ special_image_mask = input_ids == self.vocabulary_mapping.image_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ n_image_features = image_features.shape[0] * image_features.shape[1]
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ torch_compilable_check(
+ inputs_embeds[special_image_mask].numel() == image_features.numel(),
+ f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}",
+ )
+ return special_image_mask
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ image_sizes: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithPast:
+ r"""
+ image_sizes (`torch.LongTensor` of shape `(batch_size, 2)`):
+ The sizes of the images in the batch, being (height, width) for each image. Image sizes can be obtained using
+ [`AutoImageProcessor`]. See [`Emu3ImageProcessor.__call__`] for details ([]`Emu3Processor`] uses
+ [`Emu3ImageProcessor`] for processing images).
+ """
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError(
+ "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+ )
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_features = self.get_image_features(pixel_values, image_sizes).pooler_output
+ image_features = torch.cat(image_features, dim=0)
+ special_image_mask = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_features
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs = self.text_model(
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ return outputs
+
+
+class Emu3ForConditionalGeneration(Emu3PreTrainedModel, GenerationMixin):
+ output_modalities = ("image", "text")
+ _tied_weights_keys = {"lm_head.weight": "model.text_model.embed_tokens.weight"}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = Emu3Model(config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.model.set_input_embeddings(value)
+
+ def get_output_embeddings(self) -> nn.Module:
+ return self.lm_head
+
+ def decode_image_tokens(self, **kwargs):
+ return self.model.decode_image_tokens(**kwargs)
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ image_sizes: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ labels: torch.LongTensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithPast:
+ r"""
+ image_sizes (`torch.LongTensor` of shape `(batch_size, 2)`):
+ The sizes of the images in the batch, being (height, width) for each image. Image sizes can be obtained using
+ [`AutoImageProcessor`]. See [`Emu3ImageProcessor.__call__`] for details ([]`Emu3Processor`] uses
+ [`Emu3ImageProcessor`] for processing images).
+ 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]`.
+
+ Example:
+
+ ```python
+ >>> from transformers import Emu3Processor, Emu3ForConditionalGeneration
+ >>> import torch
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from PIL import Image
+
+ >>> model = Emu3ForConditionalGeneration.from_pretrained("BAAI/Emu3-Chat-hf", dtype=torch.bfloat16)
+ >>> processor = Emu3Processor.from_pretrained("BAAI/Emu3-Chat-hf")
+
+ >>> conversation = [
+ ... {
+ ... "role": "system",
+ ... "content": [
+ ... {"type": "text", "text": "You are a helpful assistant."},
+ ... ],
+ ... },
+ ... {
+ ... "role": "user",
+ ... "content": [
+ ... {"type": "image"},
+ ... {"type": "text", "text": "Please describe the image."},
+ ... ],
+ ... },
+ ... ]
+
+ >>> prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+ >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> inputs = processor(images=[image], text=[prompt], return_tensors="pt").to(model.device, torch.bfloat16)
+
+ >>> generated_ids = model.generate(**inputs, max_new_tokens=100, do_sample=False)
+ >>> processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+ ```"""
+ 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,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+ )
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ position_ids=position_ids,
+ pixel_values=pixel_values,
+ use_cache=use_cache,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ if not is_first_iteration and use_cache:
+ model_inputs["pixel_values"] = None
+
+ return model_inputs
+
+
+__all__ = [
+ "Emu3ForConditionalGeneration",
+ "Emu3ForCausalLM",
+ "Emu3TextModel",
+ "Emu3PreTrainedModel",
+ "Emu3VQVAE",
+ "Emu3Model",
+]
diff --git a/third_party/transformers/src/transformers/models/emu3/modular_emu3.py b/third_party/transformers/src/transformers/models/emu3/modular_emu3.py
new file mode 100644
index 0000000000000000000000000000000000000000..598687892727e38ba8591628e2fbe3330f9a25f0
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/emu3/modular_emu3.py
@@ -0,0 +1,1230 @@
+# Copyright 2024 HuggingFace Inc. team. All rights reserved.
+#
+#
+# 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.
+
+import math
+from dataclasses import dataclass
+from functools import cached_property
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ... import initialization as init
+from ...cache_utils import Cache
+from ...generation import GenerationMixin
+from ...modeling_outputs import BaseModelOutputWithPooling, CausalLMOutputWithPast
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import auto_docstring, can_return_tuple, logging, torch_compilable_check
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from ..chameleon.modeling_chameleon import (
+ ChameleonPreTrainedModel,
+ ChameleonVQVAEEncoderConvDownsample,
+)
+from ..llama.modeling_llama import LlamaAttention, LlamaDecoderLayer, LlamaForCausalLM, LlamaModel, TransformersKwargs
+from ..siglip.modeling_siglip import SiglipAttention
+from .configuration_emu3 import Emu3Config, Emu3TextConfig, Emu3VQVAEConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+@auto_docstring
+class Emu3VQVAEModelOutput(BaseModelOutputWithPooling):
+ r"""
+ image_tokens (`torch.LongTensor` of shape `(batch_size, config.vocab_size`):
+ Indices of the image tokens predicted by the VQ-VAE model.
+ """
+
+ image_tokens: torch.LongTensor | None = None
+
+
+class Emu3Attention(LlamaAttention):
+ pass
+
+
+# Has extra dropout which no other model in the library has
+class Emu3DecoderLayer(LlamaDecoderLayer):
+ def __init__(self, config: Emu3Config, layer_idx: int):
+ super().__init__(config, layer_idx)
+ self.dropout = nn.Dropout(config.attention_dropout)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = residual + self.dropout(hidden_states)
+
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + self.dropout(hidden_states)
+ return hidden_states
+
+
+class Emu3VQVAEVectorQuantizer(nn.Module):
+ """
+ A module for vector quantization using learned embedding vectors.
+
+ This module implements the quantization process similar to te one described in
+ the VQ-VAE (Vector Quantized Variational AutoEncoder) paper. It quantizes continuous
+ input vectors into discrete codebook vectors, which are learned during training.
+ Current implementation improves over previous ones by avoiding costly matrix multiplications
+ and allowing for post-hoc remapping of indices.
+ """
+
+ def __init__(self, config: Emu3VQVAEConfig):
+ super().__init__()
+ self.embedding = nn.Embedding(config.codebook_size, config.embed_dim)
+ self.embedding.weight.data.uniform_(-1.0 / config.codebook_size, 1.0 / config.codebook_size)
+
+ def forward(self, hidden_state: torch.Tensor):
+ batch_size, temporal, channels, height, width = hidden_state.shape
+ hidden_state = hidden_state.permute(0, 1, 3, 4, 2).contiguous()
+ hidden_state_flattened = hidden_state.view(-1, channels)
+
+ # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+ hidden_state_sum = torch.sum(hidden_state_flattened**2, dim=1, keepdim=True)
+ embedding_sum = torch.sum(self.embedding.weight**2, dim=1)
+
+ # "bd,dn->bn",
+ distances = 2 * torch.matmul(hidden_state_flattened, self.embedding.weight.transpose(0, 1))
+ distances = hidden_state_sum + embedding_sum - distances
+
+ min_encoding_indices = torch.argmin(distances, dim=1)
+ min_encoding_indices = min_encoding_indices.view(batch_size, temporal, height, width)
+ return min_encoding_indices
+
+
+class Emu3VQVAEEncoderConvDownsample(ChameleonVQVAEEncoderConvDownsample):
+ pass
+
+
+class Emu3VQVAEEncoderConvUpsample(nn.Module):
+ def __init__(self, in_channels):
+ super().__init__()
+ self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1)
+
+ def forward(self, hidden_states):
+ hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
+ hidden_states = self.conv(hidden_states)
+ return hidden_states
+
+
+class Emu3VQVAEConv3d(nn.Module):
+ def __init__(
+ self,
+ in_channel: int,
+ out_channel: int,
+ kernel_size: tuple[int],
+ stride: tuple[int],
+ ):
+ super().__init__()
+
+ padding_sizes = [one_kernel - one_stride for one_kernel, one_stride in zip(kernel_size[1:], stride[1:])]
+ self.padding = ()
+ for pad_size in padding_sizes[::-1]:
+ self.padding += (pad_size // 2 + pad_size % 2, pad_size // 2)
+ self.padding += (2, 0)
+
+ self.conv = nn.Conv3d(
+ in_channel,
+ out_channel,
+ kernel_size,
+ stride=stride,
+ )
+
+ def forward(self, hidden_states: torch.Tensor):
+ hidden_states = F.pad(hidden_states, self.padding)
+ hidden_states = self.conv(hidden_states)
+ return hidden_states
+
+
+class Emu3VQVAESpatialNorm(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ ):
+ super().__init__()
+ self.norm_layer = nn.GroupNorm(
+ num_channels=out_channels,
+ num_groups=32,
+ eps=1e-6,
+ affine=True,
+ )
+
+ self.conv_y = nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ )
+ self.conv_b = nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ )
+
+ def forward(self, hidden_states: torch.Tensor, quant_states: torch.Tensor):
+ quant_states = F.interpolate(quant_states, size=hidden_states.shape[-2:], mode="nearest")
+ hidden_states = self.norm_layer(hidden_states)
+ hidden_states = hidden_states * self.conv_y(quant_states) + self.conv_b(quant_states)
+ return hidden_states
+
+
+class Emu3VQVAETemporalUpsample(nn.Module):
+ def __init__(
+ self,
+ in_channel: int,
+ out_channel: int,
+ ):
+ super().__init__()
+ self.conv = Emu3VQVAEConv3d(
+ in_channel,
+ out_channel,
+ kernel_size=(3, 3, 3),
+ stride=(1, 1, 1),
+ )
+
+ def forward(self, hidden_states: torch.Tensor):
+ batch_size, channels, temporal, height, width = hidden_states.shape
+ hidden_states = hidden_states.permute(0, 1, 3, 4, 2).contiguous().view(batch_size, -1, temporal)
+ hidden_states = F.interpolate(hidden_states, scale_factor=2.0, mode="nearest")
+ hidden_states = hidden_states.view(batch_size, channels, height, width, -1).permute(0, 1, 4, 2, 3).contiguous()
+ hidden_states = self.conv(hidden_states)
+ return hidden_states
+
+
+class Emu3VQVAETemporalDownsample(nn.Module):
+ def __init__(
+ self,
+ in_channel: int,
+ out_channel: int,
+ ):
+ super().__init__()
+ self.conv = Emu3VQVAEConv3d(
+ in_channel,
+ out_channel,
+ kernel_size=(4, 3, 3),
+ stride=(2, 1, 1),
+ )
+
+ def forward(self, hidden_states: torch.Tensor):
+ hidden_states = self.conv(hidden_states)
+ return hidden_states
+
+
+class Emu3VQVAETemporalResnetBlock(nn.Module):
+ def __init__(
+ self,
+ in_channels,
+ out_channels=None,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ self.out_channels = in_channels if out_channels is None else out_channels
+
+ self.norm1 = nn.BatchNorm3d(in_channels)
+ self.conv1 = Emu3VQVAEConv3d(
+ in_channels,
+ out_channels,
+ kernel_size=(3, 3, 3),
+ stride=(1, 1, 1),
+ )
+ self.norm2 = nn.BatchNorm3d(out_channels)
+ self.conv2 = Emu3VQVAEConv3d(
+ out_channels,
+ out_channels,
+ kernel_size=(3, 3, 3),
+ stride=(1, 1, 1),
+ )
+ if self.in_channels != self.out_channels:
+ self.nin_shortcut = nn.Conv3d(
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ )
+
+ def forward(self, hidden_states):
+ residual = hidden_states
+ hidden_states = self.norm1(hidden_states)
+ hidden_states *= torch.sigmoid(hidden_states)
+ hidden_states = self.conv1(hidden_states)
+
+ hidden_states = self.norm2(hidden_states)
+ hidden_states *= torch.sigmoid(hidden_states)
+ hidden_states = self.conv2(hidden_states)
+
+ if self.in_channels != self.out_channels:
+ residual = self.nin_shortcut(residual)
+
+ return residual + hidden_states
+
+
+class Emu3VQVAEResnetBlock(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int | None = None,
+ quant_channels: int | None = None,
+ ):
+ super().__init__()
+ self.in_channels = in_channels
+ out_channels = in_channels if out_channels is None else out_channels
+ self.out_channels = out_channels
+ self.quant_channels = quant_channels
+
+ if quant_channels is None:
+ self.norm1 = nn.GroupNorm(num_channels=in_channels, num_groups=32, eps=1e-6, affine=True)
+ self.norm2 = nn.GroupNorm(num_channels=out_channels, num_groups=32, eps=1e-6, affine=True)
+ else:
+ self.norm1 = Emu3VQVAESpatialNorm(quant_channels, in_channels)
+ self.norm2 = Emu3VQVAESpatialNorm(quant_channels, out_channels)
+
+ self.conv1 = nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ self.conv2 = nn.Conv2d(
+ out_channels,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ if self.in_channels != self.out_channels:
+ self.nin_shortcut = nn.Conv2d(
+ in_channels,
+ out_channels,
+ kernel_size=1,
+ stride=1,
+ padding=0,
+ )
+
+ def forward(self, hidden_states: torch.Tensor, quant_channels: torch.Tensor | None = None):
+ norm_args = () if self.quant_channels is None else (quant_channels,)
+
+ residual = hidden_states
+ hidden_states = self.norm1(hidden_states, *norm_args)
+ hidden_states *= torch.sigmoid(hidden_states)
+ hidden_states = self.conv1(hidden_states)
+
+ hidden_states = self.norm2(hidden_states, *norm_args)
+ hidden_states *= torch.sigmoid(hidden_states)
+ hidden_states = self.conv2(hidden_states)
+
+ if self.in_channels != self.out_channels:
+ residual = self.nin_shortcut(residual)
+
+ return residual + hidden_states
+
+
+class Emu3VQVAEAttentionBlock(SiglipAttention):
+ def __init__(self, config: Emu3VQVAEConfig):
+ super().__init__(config)
+
+ # for compatibility with the attention interface
+ self.num_key_value_groups = 1
+
+
+class Emu3VQVAEGroupNorm(nn.GroupNorm):
+ """
+ Same as the torch GroupNorm with the only difference that this ones accepts
+ an optional kwarg `quant_states` which is not used. This class makes it easier to
+ use SpatialNorm or GroupNorm without conditionals
+ """
+
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def forward(self, input, quant_states=None):
+ return F.group_norm(input, self.num_groups, self.weight, self.bias, self.eps)
+
+
+class Emu3VQVAEMiddleBlock(nn.Module):
+ def __init__(self, config, in_channels, quant_channels=None):
+ super().__init__()
+
+ self.block_1 = Emu3VQVAEResnetBlock(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ quant_channels=quant_channels,
+ )
+ self.attn_1 = Emu3VQVAEAttentionBlock(config)
+ if quant_channels is None:
+ self.attn_norm = Emu3VQVAEGroupNorm(num_channels=in_channels, num_groups=32, eps=1e-6, affine=True)
+ else:
+ self.attn_norm = Emu3VQVAESpatialNorm(quant_channels, in_channels)
+
+ self.block_2 = Emu3VQVAEResnetBlock(
+ in_channels=in_channels,
+ out_channels=in_channels,
+ quant_channels=quant_channels,
+ )
+
+ def forward(self, hidden_states: torch.FloatTensor, quant_states: torch.FloatTensor | None = None):
+ hidden_states = self.block_1(hidden_states, quant_states)
+ residual = hidden_states
+ hidden_states = self.attn_norm(hidden_states, quant_states)
+ batch_size, channels, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)
+ hidden_states = self.attn_1(hidden_states)[0]
+ hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
+ hidden_states = residual + hidden_states
+ hidden_states = self.block_2(hidden_states, quant_states)
+ return hidden_states
+
+
+class Emu3VQVAEDownBlock(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ self.num_resolutions = len(config.channel_multiplier)
+ self.num_res_blocks = config.num_res_blocks
+ base_channels = config.base_channels
+ channel_multiplier = config.channel_multiplier
+
+ in_channel_multiplier = (1,) + tuple(channel_multiplier)
+ self.in_channel_multiplier = in_channel_multiplier
+ self.down = nn.ModuleList()
+ for i_level in range(self.num_resolutions):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ attn_norms = nn.ModuleList()
+ block_in = base_channels * in_channel_multiplier[i_level]
+ block_out = base_channels * channel_multiplier[i_level]
+ for i_block in range(self.num_res_blocks):
+ block.append(
+ Emu3VQVAEResnetBlock(
+ in_channels=block_in,
+ out_channels=block_out,
+ )
+ )
+ block_in = block_out
+ if config.attn_resolutions is not None and i_level in config.attn_resolutions:
+ attn.append(Emu3VQVAEAttentionBlock(config))
+ attn_norms.append(nn.GroupNorm(num_channels=block_in, num_groups=32, eps=1e-6, affine=True))
+
+ down = nn.Module()
+ down.block = block
+ down.attn = attn
+ down.attn_norms = attn_norms
+ if i_level != self.num_resolutions - 1:
+ down.downsample = Emu3VQVAEEncoderConvDownsample(block_in)
+ self.down.append(down)
+
+ def forward(self, hidden_states: torch.FloatTensor):
+ for i_level, blocks in enumerate(self.down):
+ for i_block in range(self.num_res_blocks):
+ hidden_states = blocks.block[i_block](hidden_states)
+ if len(blocks.attn) > 0:
+ residual = hidden_states
+ hidden_states = blocks.attn_norms[i_block](hidden_states)
+
+ batch_size, channels, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)
+ hidden_states = blocks.attn[i_block](hidden_states)[0]
+
+ hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
+ hidden_states = residual + hidden_states
+
+ if i_level != self.num_resolutions - 1:
+ hidden_states = blocks.downsample(hidden_states)
+
+ return hidden_states
+
+
+class Emu3VQVAEUpBlock(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ self.num_resolutions = len(config.channel_multiplier)
+ self.num_res_blocks = config.num_res_blocks
+
+ quant_channels = config.embed_dim
+ block_in = config.base_channels * config.channel_multiplier[-1]
+
+ self.up = nn.ModuleList()
+ for i_level in reversed(range(self.num_resolutions)):
+ block = nn.ModuleList()
+ attn = nn.ModuleList()
+ attn_norms = nn.ModuleList()
+ block_out = config.base_channels * config.channel_multiplier[i_level]
+ for i_block in range(self.num_res_blocks + 1):
+ block.append(
+ Emu3VQVAEResnetBlock(
+ in_channels=block_in,
+ out_channels=block_out,
+ quant_channels=quant_channels,
+ )
+ )
+ block_in = block_out
+ if i_level in config.attn_resolutions:
+ attn.append(Emu3VQVAEAttentionBlock(config))
+ attn_norms.append(Emu3VQVAESpatialNorm(quant_channels, block_in))
+
+ up = nn.Module()
+ up.block = block
+ up.attn = attn
+ up.attn_norms = attn_norms
+ if i_level != 0:
+ up.upsample = Emu3VQVAEEncoderConvUpsample(block_in)
+
+ self.up.insert(0, up)
+
+ def forward(self, hidden_states: torch.FloatTensor, quant_states: torch.FloatTensor):
+ for i_level, blocks in enumerate(self.up[::-1]):
+ for i_block in range(self.num_res_blocks + 1):
+ hidden_states = blocks.block[i_block](hidden_states, quant_states)
+ if len(blocks.attn) > 0:
+ residual = hidden_states
+ hidden_states = blocks.attn_norms[i_block](hidden_states, quant_states)
+
+ batch_size, channels, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channels, height * width).transpose(1, 2)
+ hidden_states = blocks.attn[i_block](hidden_states)[0]
+
+ hidden_states = hidden_states.reshape(batch_size, height, width, channels).permute(0, 3, 1, 2)
+ hidden_states = residual + hidden_states
+ if i_level != len(self.up) - 1:
+ hidden_states = blocks.upsample(hidden_states)
+
+ return hidden_states
+
+
+class Emu3VQVAEEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ base_channels = config.base_channels
+ in_channels = config.in_channels
+ double_latent = config.double_latent
+ latent_channels = config.latent_channels
+ channel_multiplier = config.channel_multiplier
+ out_channels = 2 * latent_channels if double_latent else latent_channels
+ block_in = base_channels * channel_multiplier[-1]
+
+ self.conv_in = torch.nn.Conv2d(in_channels, base_channels, kernel_size=3, stride=1, padding=1)
+ self.down_block = Emu3VQVAEDownBlock(config)
+ self.middle_block = Emu3VQVAEMiddleBlock(config, block_in)
+
+ self.norm_out = torch.nn.GroupNorm(num_groups=32, num_channels=block_in, eps=1e-6, affine=True)
+ self.conv_out = torch.nn.Conv2d(
+ block_in,
+ out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ temporal_down_blocks = int(math.log2(config.temporal_downsample_factor))
+ self.time_conv = nn.ModuleList()
+ self.time_res_stack = nn.ModuleList()
+
+ for i in range(temporal_down_blocks):
+ conv = Emu3VQVAETemporalDownsample(out_channels, out_channels)
+ self.time_conv.append(conv)
+
+ for _ in range(config.num_res_blocks):
+ time_res_conv = Emu3VQVAETemporalResnetBlock(
+ in_channels=out_channels,
+ out_channels=out_channels,
+ )
+ self.time_res_stack.append(time_res_conv)
+
+ def forward(self, pixel_values: torch.LongTensor):
+ temporal_dim = pixel_values.shape[1]
+ pixel_values = pixel_values.reshape(-1, *pixel_values.shape[2:])
+
+ # downsampling & middle
+ hidden_states = self.conv_in(pixel_values)
+ hidden_states = self.down_block(hidden_states)
+ hidden_states = self.middle_block(hidden_states)
+
+ # end
+ hidden_states = self.norm_out(hidden_states)
+ hidden_states *= torch.sigmoid(hidden_states)
+ hidden_states = self.conv_out(hidden_states)
+
+ hidden_states = hidden_states.reshape(-1, temporal_dim, *hidden_states.shape[1:])
+ hidden_states = hidden_states.permute(0, 2, 1, 3, 4)
+
+ # temporal convs
+ for conv in self.time_conv:
+ hidden_states = conv(hidden_states)
+ hidden_states *= torch.sigmoid(hidden_states)
+
+ for layer in self.time_res_stack:
+ hidden_states = layer(hidden_states)
+
+ hidden_states = hidden_states.permute(0, 2, 1, 3, 4)
+
+ return hidden_states
+
+
+class Emu3VQVAEDecoder(nn.Module):
+ def __init__(self, config: Emu3VQVAEConfig):
+ super().__init__()
+
+ quant_channels = config.embed_dim
+ block_in = config.base_channels * config.channel_multiplier[-1]
+ self.time_res_stack = nn.ModuleList()
+ for _ in range(config.num_res_blocks):
+ time_res_conv = Emu3VQVAETemporalResnetBlock(
+ in_channels=config.latent_channels, out_channels=config.latent_channels
+ )
+ self.time_res_stack.append(time_res_conv)
+
+ temp_upsample_block_num = int(math.log2(config.temporal_downsample_factor))
+ self.time_conv = nn.ModuleList()
+ for i in range(temp_upsample_block_num):
+ conv = Emu3VQVAETemporalUpsample(config.latent_channels, config.latent_channels)
+ self.time_conv.append(conv)
+
+ self.conv_in = nn.Conv2d(
+ config.latent_channels,
+ block_in,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ self.middle_block = Emu3VQVAEMiddleBlock(config, block_in, quant_channels=quant_channels)
+ self.up_block = Emu3VQVAEUpBlock(config)
+
+ block_in = config.base_channels * config.channel_multiplier[0]
+ self.norm_out = Emu3VQVAESpatialNorm(quant_channels, block_in)
+ self.conv_out = nn.Conv2d(
+ block_in,
+ config.out_channels,
+ kernel_size=3,
+ stride=1,
+ padding=1,
+ )
+
+ def forward(self, hidden_states: torch.Tensor, quant_states: torch.Tensor):
+ hidden_quant_states = torch.cat((hidden_states, quant_states), dim=0)
+ hidden_quant_states = hidden_quant_states.permute(0, 2, 1, 3, 4)
+
+ # temporal convs
+ for layer in self.time_res_stack:
+ hidden_quant_states = layer(hidden_quant_states)
+
+ for layer in self.time_conv:
+ hidden_quant_states = layer(hidden_quant_states)
+ hidden_quant_states *= torch.sigmoid(hidden_quant_states)
+
+ hidden_quant_states = hidden_quant_states.permute(0, 2, 1, 3, 4)
+ hidden_states, quant_states = torch.chunk(hidden_quant_states, 2, dim=0)
+ hidden_states = hidden_states.reshape(-1, *hidden_states.shape[2:])
+ quant_states = quant_states.reshape(-1, *quant_states.shape[2:])
+
+ hidden_states = self.conv_in(hidden_states)
+
+ # middle & upsampling
+ hidden_states = self.middle_block(hidden_states, quant_states)
+ hidden_states = self.up_block(hidden_states, quant_states)
+
+ hidden_states = self.norm_out(hidden_states, quant_states)
+ hidden_states *= torch.sigmoid(hidden_states)
+ hidden_states = self.conv_out(hidden_states)
+
+ return hidden_states
+
+
+@auto_docstring(
+ custom_intro="""
+ The VQ-VAE model used in Emu3 for encoding/decoding images into discrete tokens.
+ This model follows the "Make-a-scene: Scene-based text-to-image generation with human priors" paper from
+ [ Oran Gafni, Adam Polyak, Oron Ashual, Shelly Sheynin, Devi Parikh, and Yaniv
+ Taigman](https://huggingface.co/papers/2203.13131).
+ """
+)
+class Emu3VQVAE(PreTrainedModel):
+ config: Emu3VQVAEConfig
+ base_model_prefix = "emuvideovq"
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+ _supports_sdpa = True
+ _supports_flash_attn = True
+ _supports_flex_attn = True
+ _supports_attention_backend = True
+ _no_split_modules = [
+ "Emu3VQVAETemporalResnetBlock",
+ "Emu3VQVAEAttentionBlock",
+ "Emu3VQVAEResnetBlock",
+ "Emu3VQVAEVectorQuantizer",
+ ]
+ _can_record_outputs = {
+ "hidden_states": [Emu3VQVAEResnetBlock, Emu3VQVAETemporalResnetBlock],
+ "attentions": Emu3VQVAEAttentionBlock,
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ if isinstance(module, (nn.Conv2d, nn.Conv3d)):
+ init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu")
+ if module.bias is not None:
+ fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(module.weight)
+ bound = 1 / math.sqrt(fan_in)
+ init.uniform_(module.bias, -bound, bound)
+ elif isinstance(module, nn.Linear):
+ init.kaiming_uniform_(module.weight, a=math.sqrt(5))
+ if module.bias is not None:
+ fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(module.weight)
+ bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+ init.uniform_(module.bias, -bound, bound)
+ elif isinstance(module, (nn.BatchNorm2d, nn.BatchNorm3d, nn.GroupNorm)):
+ init.constant_(module.weight, 1.0)
+ init.constant_(module.bias, 0.0)
+ if getattr(module, "running_mean", None) is not None:
+ init.zeros_(module.running_mean)
+ init.ones_(module.running_var)
+ init.zeros_(module.num_batches_tracked)
+ elif isinstance(module, nn.Embedding):
+ init.normal_(module.weight)
+ # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag
+ if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False):
+ init.zeros_(module.weight[module.padding_idx])
+
+ def __init__(self, config: Emu3VQVAEConfig):
+ super().__init__(config)
+
+ self.config = config
+
+ self.encoder = Emu3VQVAEEncoder(config)
+ self.decoder = Emu3VQVAEDecoder(config)
+ self.quantize = Emu3VQVAEVectorQuantizer(config)
+ self.vision_spatial_factor = 2 ** (len(config.channel_multiplier) - 1)
+
+ self.quant_conv = Emu3VQVAEConv3d(
+ config.latent_channels, config.embed_dim, kernel_size=(3, 1, 1), stride=(1, 1, 1)
+ )
+ self.post_quant_conv = Emu3VQVAEConv3d(
+ config.embed_dim, config.latent_channels, kernel_size=(3, 1, 1), stride=(1, 1, 1)
+ )
+ self.spatial_scale_factor = 2 ** (len(config.channel_multiplier) - 1)
+ self.eval() # Emu3's VQ model is frozen
+
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def encode(
+ self, pixel_values: torch.Tensor, image_sizes: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
+ ) -> Emu3VQVAEModelOutput:
+ is_image = pixel_values.ndim == 4
+ if is_image:
+ temporal = self.config.temporal_downsample_factor
+ batch_size, channels, height, width = pixel_values.shape
+ pixel_values = pixel_values.unsqueeze(1).repeat(1, temporal, 1, 1, 1)
+ else:
+ batch_size, temporal, channels, height, width = pixel_values.shape
+
+ hidden_states = self.encoder(pixel_values)
+
+ # b t c h w -> b c t h w
+ conv_hidden_states = hidden_states.permute(0, 2, 1, 3, 4)
+ conv_hidden_states = self.quant_conv(conv_hidden_states)
+
+ # b c t h w -> b t c h w
+ conv_hidden_states = conv_hidden_states.permute(0, 2, 1, 3, 4)
+ codes = self.quantize(conv_hidden_states)
+
+ image_tokens = codes.squeeze(1) if is_image else codes
+
+ image_tokens = [
+ single_image[: int(size[0] / self.vision_spatial_factor), : int(size[1] / self.vision_spatial_factor)]
+ for single_image, size in zip(image_tokens, image_sizes)
+ ]
+
+ return Emu3VQVAEModelOutput(
+ last_hidden_state=hidden_states,
+ image_tokens=image_tokens,
+ )
+
+ def decode(self, hidden_states: torch.Tensor):
+ is_image = hidden_states.ndim == 3
+ if is_image:
+ hidden_states = hidden_states.unsqueeze(1)
+
+ batch_size, temporal, height, width = hidden_states.shape
+ quant = self.quantize.embedding(hidden_states.flatten())
+
+ channels = quant.shape[-1]
+ quant = quant.view(batch_size, temporal, height, width, channels).permute(0, 4, 1, 2, 3).contiguous()
+ post_quant = self.post_quant_conv(quant)
+
+ quant = quant.permute(0, 2, 1, 3, 4)
+ post_quant = post_quant.permute(0, 2, 1, 3, 4)
+
+ video = self.decoder(post_quant, quant)
+ video = video.reshape(
+ batch_size,
+ temporal * self.config.temporal_downsample_factor,
+ self.config.out_channels,
+ height * self.spatial_scale_factor,
+ width * self.spatial_scale_factor,
+ )
+ return video[:, 0] if is_image else video
+
+
+class Emu3ImageVocabularyMapping:
+ """
+ A class for mapping discrete image tokens from VQGAN to BPE tokens.
+ """
+
+ def __init__(self, vocab_map):
+ self.vocab_map = vocab_map
+ self.eol_token_id = vocab_map.get("<|extra_200|>")
+ self.image_token_id = vocab_map.get("")
+
+ @cached_property
+ def image_tokens(self):
+ return sorted([val for name, val in self.vocab_map.items() if name.startswith("<|visual token")])
+
+ @cached_property
+ def image_tokens_str(self):
+ return sorted([name for name, val in self.vocab_map.items() if name.startswith("<|visual token")])
+
+ @cached_property
+ def img2bpe(self):
+ return {int(token[-8:-2]): self.vocab_map[token] for token in self.image_tokens_str}
+
+ @cached_property
+ def bpe2img(self):
+ return {v: k for k, v in self.img2bpe.items()}
+
+ @cached_property
+ def bpe2img_mapping_tensor(self):
+ mapping = torch.zeros(max(self.bpe2img.keys()) + 1, dtype=torch.int)
+ for k, v in self.bpe2img.items():
+ mapping[k] = v
+ return mapping
+
+ @cached_property
+ def img2bpe_mapping_tensor(self):
+ mapping = torch.zeros(max(self.img2bpe.keys()) + 1, dtype=torch.int)
+ for k, v in self.img2bpe.items():
+ mapping[k] = v
+ return mapping
+
+ def convert_img2bpe(self, img_batch: list[torch.Tensor]) -> torch.Tensor:
+ device = img_batch.device
+ eol_row = torch.ones((img_batch.shape[0], 1), dtype=torch.int) * self.eol_token_id
+ img_tokens = self.img2bpe_mapping_tensor[img_batch.to("cpu")]
+ img_tokens = torch.cat([img_tokens, eol_row], dim=-1)
+ return img_tokens.to(device)
+
+ def convert_bpe2img(self, img_batch: torch.Tensor) -> torch.Tensor:
+ device = img_batch.device
+ img_batch = img_batch[..., :-1] # remove last row of EOL tokens
+ img_tokens = self.bpe2img_mapping_tensor[img_batch.to("cpu")]
+ return img_tokens.to(device)
+
+
+class Emu3PreTrainedModel(ChameleonPreTrainedModel):
+ _no_split_modules = [
+ "Emu3DecoderLayer",
+ ]
+ _supports_flex_attn = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": Emu3DecoderLayer,
+ "attentions": Emu3Attention,
+ }
+
+
+class Emu3TextModel(LlamaModel, Emu3PreTrainedModel):
+ config: Emu3TextConfig
+
+ def __init__(self, config: Emu3TextConfig):
+ super().__init__(config)
+ self.layers = nn.ModuleList(
+ [Emu3DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+
+
+class Emu3ForCausalLM(LlamaForCausalLM, Emu3PreTrainedModel, GenerationMixin):
+ config: Emu3TextConfig
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = Emu3TextModel(config)
+
+ def forward(**super_kwargs):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import Emu3Processor, Emu3ForConditionalGeneration
+ >>> import torch
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from PIL import Image
+
+ >>> model = Emu3ForCausalLM.from_pretrained("BAAI/Emu3-Chat-hf", dtype=torch.bfloat16)
+ >>> processor = Emu3Processor.from_pretrained("BAAI/Emu3-Chat-hf")
+
+ >>> inputs = processor(text=["Can you write me a poem about winter."], return_tensors="pt").to(model.device)
+
+ >>> generated_ids = model.generate(**inputs, max_new_tokens=100, do_sample=False)
+ >>> processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+ ```"""
+ super().forward()
+
+
+class Emu3Model(Emu3PreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.text_model = Emu3TextModel._from_config(config.text_config)
+ self.vqmodel = Emu3VQVAE(config.vq_config)
+ self.vocabulary_mapping = Emu3ImageVocabularyMapping(config.vocabulary_map)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.text_model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.text_model.set_input_embeddings(value)
+
+ def get_image_tokens(self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor) -> torch.LongTensor:
+ """
+ Tokenizes images into discrete tokens with VQGAN module. Converts
+ obtained image tokens into BPE tokens and wraps with "boi" and "eoi"
+ special tokens.
+
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_sizes (`torch.LongTensor` of shape `(batch_size, 2)`):
+ The sizes of the images in the batch, being (height, width) for each image.
+ """
+ vqmodel_outputs: Emu3VQVAEModelOutput = self.vqmodel.encode(pixel_values, image_sizes, return_dict=True)
+ bpe_tokens_list = [
+ self.vocabulary_mapping.convert_img2bpe(tokens).flatten() for tokens in vqmodel_outputs.image_tokens
+ ]
+ bpe_tokens = torch.cat(bpe_tokens_list)
+ return bpe_tokens
+
+ @can_return_tuple
+ @auto_docstring(
+ custom_intro="Tokenizes images into discrete tokens with VQGAN module and embeds them with text embeddings layer"
+ )
+ def get_image_features(
+ self, pixel_values: torch.FloatTensor, image_sizes: torch.LongTensor, **kwargs: Unpack[TransformersKwargs]
+ ) -> tuple | Emu3VQVAEModelOutput:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)):
+ The tensors corresponding to the input images.
+ """
+ vqmodel_outputs: Emu3VQVAEModelOutput = self.vqmodel.encode(
+ pixel_values, image_sizes, return_dict=True, **kwargs
+ )
+ split_sizes = [
+ (height // self.vqmodel.vision_spatial_factor) * (width // self.vqmodel.vision_spatial_factor + 1)
+ for height, width in image_sizes
+ ]
+ bpe_tokens_list = [
+ self.vocabulary_mapping.convert_img2bpe(tokens).flatten() for tokens in vqmodel_outputs.image_tokens
+ ]
+ bpe_tokens = torch.cat(bpe_tokens_list)
+ image_embeddings = self.get_input_embeddings()(bpe_tokens)
+ image_features = torch.split(image_embeddings, split_sizes)
+ vqmodel_outputs.pooler_output = image_features
+
+ return vqmodel_outputs
+
+ @torch.no_grad()
+ def decode_image_tokens(self, image_tokens: torch.LongTensor, height: int, width: int):
+ """
+ Decodes generated image tokens from language model to continuous pixel values
+ with VQGAN module via upsampling.
+
+ Args:
+ image_tokens (`torch.LongTensor` of shape `(batch_size, num_of_tokens)`):
+ The tensors corresponding to the input images.
+ height (`int`):
+ Height of the generated image before upsampling.
+ width (`int`):
+ Width of the generated image before upsampling.
+ """
+ sequences = image_tokens[:, :-3].view(-1, height, width + 1)
+ image_tokens = self.vocabulary_mapping.convert_bpe2img(sequences)
+ image = self.vqmodel.decode(image_tokens)
+ return image
+
+ def get_placeholder_mask(
+ self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor
+ ):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+ equal to the length of multimodal features. If the lengths are different, an error is raised.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.vocabulary_mapping.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ else:
+ special_image_mask = input_ids == self.vocabulary_mapping.image_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ n_image_features = image_features.shape[0] * image_features.shape[1]
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ torch_compilable_check(
+ inputs_embeds[special_image_mask].numel() == image_features.numel(),
+ f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}",
+ )
+ return special_image_mask
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ image_sizes: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithPast:
+ r"""
+ image_sizes (`torch.LongTensor` of shape `(batch_size, 2)`):
+ The sizes of the images in the batch, being (height, width) for each image. Image sizes can be obtained using
+ [`AutoImageProcessor`]. See [`Emu3ImageProcessor.__call__`] for details ([]`Emu3Processor`] uses
+ [`Emu3ImageProcessor`] for processing images).
+ """
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError(
+ "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
+ )
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_features = self.get_image_features(pixel_values, image_sizes).pooler_output
+ image_features = torch.cat(image_features, dim=0)
+ special_image_mask = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_features
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
+ outputs = self.text_model(
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ return outputs
+
+
+class Emu3ForConditionalGeneration(Emu3PreTrainedModel, GenerationMixin):
+ output_modalities = ("image", "text")
+ _tied_weights_keys = {"lm_head.weight": "model.text_model.embed_tokens.weight"}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = Emu3Model(config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.model.set_input_embeddings(value)
+
+ def get_output_embeddings(self) -> nn.Module:
+ return self.lm_head
+
+ def decode_image_tokens(self, **kwargs):
+ return self.model.decode_image_tokens(**kwargs)
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ image_sizes: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ labels: torch.LongTensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithPast:
+ r"""
+ image_sizes (`torch.LongTensor` of shape `(batch_size, 2)`):
+ The sizes of the images in the batch, being (height, width) for each image. Image sizes can be obtained using
+ [`AutoImageProcessor`]. See [`Emu3ImageProcessor.__call__`] for details ([]`Emu3Processor`] uses
+ [`Emu3ImageProcessor`] for processing images).
+ 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]`.
+
+ Example:
+
+ ```python
+ >>> from transformers import Emu3Processor, Emu3ForConditionalGeneration
+ >>> import torch
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from PIL import Image
+
+ >>> model = Emu3ForConditionalGeneration.from_pretrained("BAAI/Emu3-Chat-hf", dtype=torch.bfloat16)
+ >>> processor = Emu3Processor.from_pretrained("BAAI/Emu3-Chat-hf")
+
+ >>> conversation = [
+ ... {
+ ... "role": "system",
+ ... "content": [
+ ... {"type": "text", "text": "You are a helpful assistant."},
+ ... ],
+ ... },
+ ... {
+ ... "role": "user",
+ ... "content": [
+ ... {"type": "image"},
+ ... {"type": "text", "text": "Please describe the image."},
+ ... ],
+ ... },
+ ... ]
+
+ >>> prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+ >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> inputs = processor(images=[image], text=[prompt], return_tensors="pt").to(model.device, torch.bfloat16)
+
+ >>> generated_ids = model.generate(**inputs, max_new_tokens=100, do_sample=False)
+ >>> processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+ ```"""
+ 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,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+ )
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ position_ids=position_ids,
+ pixel_values=pixel_values,
+ use_cache=use_cache,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ if not is_first_iteration and use_cache:
+ model_inputs["pixel_values"] = None
+
+ return model_inputs
+
+
+__all__ = [
+ "Emu3ForConditionalGeneration",
+ "Emu3ForCausalLM",
+ "Emu3TextModel",
+ "Emu3PreTrainedModel",
+ "Emu3VQVAE",
+ "Emu3Model",
+]
diff --git a/third_party/transformers/src/transformers/models/emu3/processing_emu3.py b/third_party/transformers/src/transformers/models/emu3/processing_emu3.py
new file mode 100644
index 0000000000000000000000000000000000000000..ec51a6305f0bf8ef70a0f38bb5ed748120011848
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/emu3/processing_emu3.py
@@ -0,0 +1,238 @@
+# Copyright 2024 HuggingFace Inc. team. All rights reserved.
+#
+#
+# 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.
+
+
+from ...image_processing_utils import BatchFeature
+from ...image_utils import ImageInput
+from ...processing_utils import MultiModalData, ProcessingKwargs, ProcessorMixin, TextKwargs, Unpack
+from ...tokenization_utils_base import PreTokenizedInput, TextInput
+from ...utils import auto_docstring, is_vision_available
+from ...utils.import_utils import requires
+
+
+if is_vision_available():
+ from .image_processing_emu3 import Emu3ImageProcessorKwargs, smart_resize
+
+
+class Emu3TextKwargs(TextKwargs, total=False):
+ """
+ return_for_image_generation (`bool`, *optional*, defaults to `False`):
+ Whether the processed text is intended for image generation tasks. When `True`, the processor prepares
+ inputs for image generation by appending image start tokens and size information to the prompt, and
+ images should not be provided. When `False`, the processor prepares inputs for text generation from
+ images and text, requiring both inputs to be provided.
+ """
+
+ return_for_image_generation: bool
+
+
+class Emu3ProcessorKwargs(ProcessingKwargs, total=False):
+ text_kwargs: Emu3TextKwargs
+ images_kwargs: Emu3ImageProcessorKwargs
+ _defaults = {
+ "text_kwargs": {
+ "return_for_image_generation": False,
+ "return_mm_token_type_ids": False,
+ },
+ "images_kwargs": {
+ "ratio": "1:1",
+ "image_area": 518400,
+ },
+ }
+
+
+@auto_docstring
+@requires(backends=("vision",))
+class Emu3Processor(ProcessorMixin):
+ def __init__(
+ self,
+ image_processor,
+ tokenizer,
+ chat_template=None,
+ **kwargs,
+ ):
+ self.image_token = tokenizer.image_token # image_token as placeholder to be replaced by vq-vae tokens
+ self.image_token_id = tokenizer.image_token_id
+ self.image_start_token = tokenizer.boi_token # "<|image start|>" fixed tokens for start and end of image
+ self.image_end_token = tokenizer.eoi_token # "<|image end|>"
+ self.fake_token_around_image = tokenizer.image_wrapper_token # "<|image token|>" every image starts with it
+ self.eof_token = tokenizer.eof_token # "<|extra_201|>"
+ self.bos_token = tokenizer.bos_token
+ self.downsample_ratio = 8
+ super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+ @auto_docstring
+ def __call__(
+ self,
+ images: ImageInput | None = None,
+ text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] | None = None,
+ **kwargs: Unpack[Emu3ProcessorKwargs],
+ ) -> BatchFeature:
+ r"""
+ Returns:
+ [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+ - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+ - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+ `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+ `None`).
+ - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+ """
+ # check if images and text inputs are reversed for BC
+
+ if isinstance(text, str):
+ text = [text]
+ elif not isinstance(text, list) and not isinstance(text[0], str):
+ raise TypeError("Invalid input text. Please provide a string, or a list of strings")
+
+ output_kwargs = self._merge_kwargs(
+ Emu3ProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+ return_for_image_generation = output_kwargs["text_kwargs"].pop("return_for_image_generation", False)
+ ratio = output_kwargs["images_kwargs"].pop("ratio", None)
+ image_area = output_kwargs["images_kwargs"].pop("image_area", None)
+
+ if return_for_image_generation and images is not None:
+ raise ValueError("You should not provide `images` when `return_for_image_generation=True`")
+
+ if not return_for_image_generation and text is None and images is None:
+ raise ValueError("You must provide either text or images when `return_for_image_generation=False`")
+
+ image_features = {}
+ image_start_tokens = f"{self.image_start_token}"
+ image_end_tokens = f"{self.eof_token}{self.image_end_token}"
+
+ # generate text from image + text input, so we add placeholders for image tokens
+ if not return_for_image_generation and images is not None:
+ image_features = self.image_processor(images, **output_kwargs["images_kwargs"])
+ image_sizes = iter(image_features.image_sizes)
+
+ prompt_strings = []
+ for sample in text:
+ while self.image_token in sample:
+ image_size = next(image_sizes)
+ height, width = image_size
+ height = height // self.downsample_ratio
+ width = width // self.downsample_ratio
+ image_seq_length = height * (width + 1) # +1 for extra row when converting to BPE in modeling code
+
+ image_placeholder = f"{image_start_tokens}{height}*{width}{self.fake_token_around_image}{'' * image_seq_length}{image_end_tokens}"
+ sample = sample.replace(self.image_token, image_placeholder, 1)
+ sample = f"{self.bos_token}{sample}" # add BOS because GPT tokenizer doesn't add it
+ prompt_strings.append(sample)
+ text = [sample.replace("", self.image_token) for sample in prompt_strings]
+
+ # generate image from text input, so we add begin-of-image tokens from where image generation starts
+ elif return_for_image_generation:
+ height, width = self.calculate_generate_size(ratio, image_area, self.downsample_ratio)
+ image_prompt = f"{image_start_tokens}{height}*{width}{self.fake_token_around_image}"
+ text = [f"{self.bos_token}{sample}{image_prompt}" for sample in text]
+ image_features["image_sizes"] = [[height, width]] * len(text)
+
+ # else just generate from text-only input, and we do no special treatment for text
+ return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
+ return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False)
+ text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"], return_tensors=None)
+ self._check_special_mm_tokens(text, text_inputs, modalities=["image"])
+
+ if return_mm_token_type_ids:
+ text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"])
+ return BatchFeature(data={**text_inputs, **image_features}, tensor_type=return_tensors)
+
+ def _get_num_multimodal_tokens(self, image_sizes=None, **kwargs):
+ """
+ Computes the number of placeholder tokens needed for multimodal inputs with the given sizes.
+
+ Args:
+ image_sizes (`list[list[int]]`, *optional*):
+ The input sizes formatted as (height, width) per each image.
+
+ Returns:
+ `MultiModalData`: A `MultiModalData` object holding number of tokens per each of the provided
+ input modalities, along with other useful data.
+ """
+
+ vision_data = {}
+ if image_sizes is not None:
+ num_image_tokens = []
+ for height, width in image_sizes:
+ height, width = smart_resize(
+ height,
+ width,
+ self.image_processor.spatial_factor,
+ self.image_processor.min_pixels,
+ self.image_processor.max_pixels,
+ )
+ height = height // self.downsample_ratio
+ width = width // self.downsample_ratio
+ image_seq_length = height * (width + 1) # +1 for extra row when converting to BPE in modeling code
+ num_image_tokens.append(image_seq_length)
+
+ num_image_patches = [1] * len(image_sizes)
+ vision_data.update({"num_image_tokens": num_image_tokens, "num_image_patches": num_image_patches})
+
+ return MultiModalData(**vision_data)
+
+ def calculate_generate_size(self, ratio, image_area, spatial_factor):
+ width, height = map(int, ratio.split(":"))
+ current_area = width * height
+ target_ratio = (image_area / current_area) ** 0.5
+
+ token_height = int(round(height * target_ratio / spatial_factor))
+ token_width = int(round(width * target_ratio / spatial_factor))
+ return token_height, token_width
+
+ def postprocess(self, images: ImageInput, **kwargs):
+ return self.image_processor.postprocess(images, **kwargs)
+
+ def post_process_multimodal_output(
+ self, generated_outputs, skip_special_tokens=True, generation_mode=None, **kwargs
+ ):
+ """
+ Post-process the output of a multimodal model to return the requested modality output.
+ If the model cannot generated the requested modality, an error will be raised.
+
+ Args:
+ generated_outputs (`torch.Tensor` or `np.ndarray`):
+ The output of the model `generate` function. The output is expected to be a tensor of shape `(batch_size, sequence_length)`
+ or `(sequence_length,)`.
+ skip_special_tokens (`bool`, *optional*, defaults to `True`):
+ Whether or not to remove special tokens in the output. Argument passed to the tokenizer's `batch_decode` method.
+ generation_mode (`str`, *optional*):
+ Generation mode indicated which modality to output and can be one of `["text", "image", "audio"]`.
+ **kwargs:
+ Additional arguments to be passed to the tokenizer's `batch_decode method`.
+
+ Returns:
+ `list[Union[str, PIL.Image.Image]]`: The decoded text or generated image.
+ """
+ if generation_mode is None or generation_mode == "text":
+ return self.post_process_image_text_to_text(
+ generated_outputs, skip_special_tokens=skip_special_tokens, **kwargs
+ )
+
+ elif generation_mode == "image":
+ images = self.postprocess(generated_outputs, return_tensors="PIL.Image.Image")
+ return images["pixel_values"]
+
+ else:
+ raise ValueError(
+ f"{self.__class__.__name__} got an unexpected generation_mode={generation_mode}. Supported options are only `text` and `image"
+ )
+
+
+__all__ = ["Emu3Processor"]
diff --git a/third_party/transformers/src/transformers/models/eomt/__init__.py b/third_party/transformers/src/transformers/models/eomt/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..1fe617f6b8276fe99dd3a4b8a7cca32e48ce3302
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/eomt/__init__.py
@@ -0,0 +1,29 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_eomt import *
+ from .image_processing_eomt import *
+ from .image_processing_pil_eomt import *
+ from .modeling_eomt import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/eomt/configuration_eomt.py b/third_party/transformers/src/transformers/models/eomt/configuration_eomt.py
new file mode 100644
index 0000000000000000000000000000000000000000..3d81e64b0ed599b72fd8588ba985043465131510
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/eomt/configuration_eomt.py
@@ -0,0 +1,100 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/eomt/modular_eomt.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_eomt.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 Mobile Perception Systems Lab at TU/e and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="tue-mps/coco_panoptic_eomt_large_640")
+@strict
+class EomtConfig(PreTrainedConfig):
+ r"""
+ layerscale_value (`float`, *optional*, defaults to 1.0):
+ Initial value for the LayerScale parameter.
+ num_upscale_blocks (`int`, *optional*, defaults to 2):
+ Number of upsampling blocks used in the decoder or segmentation head.
+ use_swiglu_ffn (`bool`, *optional*, defaults to `False`):
+ Whether to use the SwiGLU feedforward neural network.
+ num_blocks (`int`, *optional*, defaults to 4):
+ Number of feature blocks or stages in the architecture.
+ no_object_weight (`float`, *optional*, defaults to 0.1):
+ Loss weight for the 'no object' class in panoptic/instance segmentation.
+ class_weight (`float`, *optional*, defaults to 2.0):
+ Loss weight for classification targets.
+ mask_weight (`float`, *optional*, defaults to 5.0):
+ Loss weight for mask prediction.
+ train_num_points (`int`, *optional*, defaults to 12544):
+ Number of points to sample for mask loss computation during training.
+ oversample_ratio (`float`, *optional*, defaults to 3.0):
+ Oversampling ratio used in point sampling for mask training.
+ importance_sample_ratio (`float`, *optional*, defaults to 0.75):
+ Ratio of points to sample based on importance during training.
+ num_queries (`int`, *optional*, defaults to 200):
+ Number of object queries in the Transformer.
+ num_register_tokens (`int`, *optional*, defaults to 4):
+ Number of learnable register tokens added to the transformer input.
+
+ Example:
+
+ ```python
+ >>> from transformers import EomtConfig, EomtForUniversalSegmentation
+
+ >>> # Initialize configuration
+ >>> config = EomtConfig()
+
+ >>> # Initialize model
+ >>> model = EomtForUniversalSegmentation(config)
+
+ >>> # Access config
+ >>> config = model.config
+ ```"""
+
+ model_type = "eomt"
+
+ hidden_size: int = 1024
+ num_hidden_layers: int = 24
+ num_attention_heads: int = 16
+ hidden_act: str = "gelu"
+ hidden_dropout_prob: float | int = 0.0
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-6
+ image_size: int | list[int] | tuple[int, int] = 640
+ patch_size: int | list[int] | tuple[int, int] = 16
+ num_channels: int = 3
+ mlp_ratio: int = 4
+ layerscale_value: float = 1.0
+ drop_path_rate: float | int = 0.0
+ num_upscale_blocks: int = 2
+ attention_dropout: float | int = 0.0
+ use_swiglu_ffn: bool = False
+ num_blocks: int = 4
+ no_object_weight: float = 0.1
+ class_weight: float = 2.0
+ mask_weight: float = 5.0
+ dice_weight: float = 5.0
+ train_num_points: int = 12544
+ oversample_ratio: float = 3.0
+ importance_sample_ratio: float = 0.75
+ num_queries: int = 200
+ num_register_tokens: int = 4
+
+
+__all__ = ["EomtConfig"]
diff --git a/third_party/transformers/src/transformers/models/eomt/convert_eomt_to_hf.py b/third_party/transformers/src/transformers/models/eomt/convert_eomt_to_hf.py
new file mode 100644
index 0000000000000000000000000000000000000000..1aefb9f0e62c6133d509d1516f3bf30ac5f00e55
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/eomt/convert_eomt_to_hf.py
@@ -0,0 +1,330 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+import argparse
+import gc
+import json
+import os
+import re
+
+import torch
+from huggingface_hub import snapshot_download
+
+from transformers import EomtConfig, EomtForUniversalSegmentation, EomtImageProcessorFast
+
+
+# fmt: off
+MAPPINGS = {
+ # Embeddings
+ r"network.encoder.backbone.cls_token" : r"embeddings.cls_token",
+ r"network.encoder.backbone.reg_token" : r"embeddings.register_tokens",
+ r"network.encoder.backbone.pos_embed" : r"embeddings.position_embeddings.weight",
+ r"network.encoder.backbone.patch_embed.proj" : r"embeddings.patch_embeddings.projection",
+
+ # Encoder Block
+ r"network.encoder.backbone.blocks.(\d+).norm1" : r"layers.\1.norm1",
+ r"network.encoder.backbone.blocks.(\d+).attn.proj" : r"layers.\1.attention.out_proj",
+ r"network.encoder.backbone.blocks.(\d+).ls1.gamma" : r"layers.\1.layer_scale1.lambda1",
+ r"network.encoder.backbone.blocks.(\d+).norm2" : r"layers.\1.norm2",
+ r"network.encoder.backbone.blocks.(\d+).ls2.gamma" : r"layers.\1.layer_scale2.lambda1",
+ r"network.encoder.backbone.blocks.(\d+).attn" : r"layers.\1.attention",
+
+ # Others
+ r"network.q.weight" : r"query.weight",
+ r"network.class_head" : r"class_predictor",
+ r"network.upscale.(\d+).conv1" : r"upscale_block.block.\1.conv1",
+ r"network.upscale.(\d+).conv2" : r"upscale_block.block.\1.conv2",
+ r"network.upscale.(\d+).norm" : r"upscale_block.block.\1.layernorm2d",
+ r"network.mask_head.0" : r"mask_head.fc1",
+ r"network.mask_head.2" : r"mask_head.fc2",
+ r"network.mask_head.4" : r"mask_head.fc3",
+ r"network.encoder.backbone.norm" : r"layernorm",
+ r"network.attn_mask_probs" : r"attn_mask_probs",
+}
+# fmt: on
+
+# Mappings for MLP layers, depending on the type of MLP used in ckpts.
+MLP_MAPPINGS = {
+ "swiglu_ffn": {
+ r"network.encoder.backbone.blocks.(\d+).mlp.fc1": r"layers.\1.mlp.weights_in",
+ r"network.encoder.backbone.blocks.(\d+).mlp.fc2": r"layers.\1.mlp.weights_out",
+ },
+ "vanilla_mlp": {
+ r"network.encoder.backbone.blocks.(\d+).mlp": r"layers.\1.mlp",
+ },
+}
+
+
+def convert_old_keys_to_new_keys(state_dict):
+ keys_as_text = "\n".join(state_dict.keys())
+ new_keys_as_text = keys_as_text
+ for old, repl in MAPPINGS.items():
+ if repl is None:
+ new_keys_as_text = re.sub(old, "", new_keys_as_text)
+ else:
+ new_keys_as_text = re.sub(old, repl, new_keys_as_text)
+ output_dict = dict(zip(keys_as_text.split("\n"), new_keys_as_text.split("\n")))
+ return output_dict
+
+
+def split_qkv_tensor(key, tensor):
+ """Splits a qkv tensor into separate q, k, v tensors and updates the key accordingly."""
+
+ new_keys = ["q_proj", "k_proj", "v_proj"]
+ split_size = tensor.shape[0] // 3
+ split_tensors = torch.split(tensor, split_size, dim=0)
+
+ return {key.replace("qkv", new_key): split_tensors[i] for i, new_key in enumerate(new_keys)}
+
+
+def convert_state_dict_to_hf(state_dict):
+ """Convert state dict keys to HF format."""
+ conversion_dict = convert_old_keys_to_new_keys(state_dict)
+ converted_state_dict = {}
+
+ for old_key, new_key in conversion_dict.items():
+ if new_key:
+ if "qkv" in new_key: # Detect merged attention keys and split them.
+ qkv_split_dict = split_qkv_tensor(new_key, state_dict[old_key])
+ converted_state_dict.update(qkv_split_dict)
+ else:
+ converted_state_dict[new_key] = state_dict[old_key]
+
+ for i in [
+ "network.encoder.pixel_mean",
+ "network.encoder.pixel_std",
+ ]:
+ converted_state_dict.pop(i)
+
+ # Embeddings will not have initial dimension
+ pos_embed_key = "embeddings.position_embeddings.weight"
+ converted_state_dict[pos_embed_key] = converted_state_dict[pos_embed_key].squeeze(0)
+
+ return converted_state_dict
+
+
+def ensure_model_downloaded(
+ repo_id: str | None = None, revision: str | None = None, local_dir: str | None = None
+) -> str:
+ """
+ Ensures model files are downloaded locally, downloads them if not.
+ Returns path to local files.
+
+ Args:
+ repo_id: The Hugging Face model repo ID (required if local_dir not provided)
+ revision: Optional git revision to use
+ local_dir: Optional local directory path where model files should be stored/found
+ """
+ if local_dir is not None:
+ if os.path.exists(local_dir):
+ print(f"Using provided local directory: {local_dir}")
+ else:
+ # Create the local directory if it doesn't exist
+ os.makedirs(local_dir, exist_ok=True)
+ print(f"Created local directory: {local_dir}")
+
+ if repo_id is None:
+ raise ValueError("Either repo_id or local_dir must be provided")
+
+ print(f"Ensuring {repo_id} (revision: {revision or 'latest'}) is downloaded...")
+
+ try:
+ # First try to find files locally
+ download_dir = snapshot_download(repo_id, revision=revision, local_files_only=True, local_dir=local_dir)
+ print(f"Found model files locally at {download_dir}")
+ return download_dir
+ except Exception:
+ # If files not found locally, download them
+ print(f"Downloading model files for {repo_id}...")
+ download_dir = snapshot_download(repo_id, revision=revision, local_files_only=False, local_dir=local_dir)
+ print(f"Downloaded model files to {download_dir}")
+ return download_dir
+
+
+def load_model_state_dict(input_path: str) -> dict:
+ """
+ Load model state dict, handling both single and sharded files.
+ """
+ index_path = os.path.join(input_path, "pytorch_model.bin.index.json")
+ single_file_path = os.path.join(input_path, "pytorch_model.bin")
+
+ # Check if we have a sharded model
+ if os.path.exists(index_path):
+ print("Loading sharded model...")
+ state_dict = {}
+ with open(index_path, "r") as f:
+ index = json.load(f)
+
+ # Get unique shard files and load each one only once
+ unique_shard_files = sorted(set(index["weight_map"].values()))
+ for shard_file in unique_shard_files:
+ print(f"Loading shard {shard_file}...")
+ shard_path = os.path.join(input_path, shard_file)
+ shard_dict = torch.load(shard_path, map_location="cpu")
+ state_dict.update(shard_dict)
+
+ return state_dict
+
+ # Single file model
+ elif os.path.exists(single_file_path):
+ print("Loading single file model...")
+ return torch.load(single_file_path, map_location="cpu")
+
+ else:
+ raise ValueError(f"No model files found in {input_path}")
+
+
+def convert_model(
+ repo_id=None,
+ local_dir=None,
+ output_dir=None,
+ output_hub_path=None,
+ revision=None,
+):
+ """Convert and save the model weights, processor, and configuration."""
+ if output_dir is None and output_hub_path is None:
+ raise ValueError("At least one of output_dir or output_hub_path must be specified")
+
+ if repo_id is None and local_dir is None:
+ raise ValueError("Either repo_id or local_dir must be specified")
+
+ # Create output directory if specified
+ if output_dir:
+ os.makedirs(output_dir, exist_ok=True)
+ print(f"Created/verified output directory: {output_dir}")
+
+ torch.set_default_dtype(torch.float16)
+
+ # Download or locate model files
+ input_path = ensure_model_downloaded(repo_id=repo_id, revision=revision, local_dir=local_dir)
+
+ with open(os.path.join(input_path, "config.json"), "r") as f:
+ config_data = json.load(f)
+ # Pop off unwanted keys
+ _ = config_data.pop("backbone", None)
+
+ config = EomtConfig(
+ **{
+ **config_data,
+ "layerscale_value": 1e-5,
+ }
+ )
+
+ if "semantic" in repo_id.split("_"):
+ size = {"shortest_edge": config.image_size, "longest_edge": None}
+ do_split_image = True
+ do_pad = False
+ else:
+ size = {"shortest_edge": config.image_size, "longest_edge": config.image_size}
+ do_split_image = False
+ do_pad = True
+
+ if "giant" in repo_id.split("_"):
+ config.use_swiglu_ffn = True
+ config.hidden_size = 1536
+ config.num_hidden_layers = 40
+ config.num_attention_heads = 24
+ # Update MAPPINGS for ckpts depending on the MLP type
+ MAPPINGS.update(MLP_MAPPINGS["swiglu_ffn"])
+ else:
+ MAPPINGS.update(MLP_MAPPINGS["vanilla_mlp"])
+
+ processor = EomtImageProcessorFast(size=size, do_split_image=do_split_image, do_pad=do_pad)
+
+ # Save the config and processor
+ if output_dir:
+ config.save_pretrained(output_dir)
+ processor.save_pretrained(output_dir)
+ if output_hub_path:
+ config.push_to_hub(output_hub_path)
+ processor.push_to_hub(output_hub_path)
+
+ # Initialize model with empty weights
+ print("Creating empty model...")
+ with torch.device("meta"):
+ model = EomtForUniversalSegmentation(config)
+
+ # Load and convert state dict
+ print("Loading state dict...")
+ state_dict = load_model_state_dict(input_path)
+ state_dict = convert_state_dict_to_hf(state_dict)
+
+ # Load converted state dict
+ print("Loading converted weights into model...")
+ model.load_state_dict(state_dict, strict=True, assign=True)
+
+ # Save the model
+ if output_dir:
+ print(f"Saving model to {output_dir}...")
+ model.save_pretrained(output_dir)
+ if output_hub_path:
+ print(f"Pushing model to hub at {output_hub_path}...")
+ model.push_to_hub(output_hub_path)
+
+ del state_dict, model
+ gc.collect()
+
+ # Validate the saved model if saved locally
+ if output_dir:
+ print("Reloading the local model to check if it's saved correctly...")
+ EomtForUniversalSegmentation.from_pretrained(output_dir, device_map="auto")
+ print("Local model reloaded successfully.")
+
+
+def main():
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--hf_repo_id",
+ help="HuggingFace Hub repo ID for the model",
+ default=None,
+ )
+ parser.add_argument(
+ "--local_dir",
+ help="Local directory containing the model files",
+ default=None,
+ )
+ parser.add_argument(
+ "--revision",
+ help="Specific revision to download from the Hub",
+ default=None,
+ )
+ parser.add_argument(
+ "--output_dir",
+ help="Location to write HF model locally",
+ default=None,
+ )
+ parser.add_argument(
+ "--output_hub_path",
+ help="Repository ID to push model to hub (e.g. 'username/model-name')",
+ default=None,
+ )
+ args = parser.parse_args()
+
+ if args.output_dir is None and args.output_hub_path is None:
+ raise ValueError("At least one of --output_dir or --output_hub_path must be specified")
+
+ if args.hf_repo_id is None and args.local_dir is None:
+ raise ValueError("Either --hf_repo_id or --local_dir must be specified")
+
+ convert_model(
+ repo_id=args.hf_repo_id,
+ local_dir=args.local_dir,
+ output_dir=args.output_dir,
+ output_hub_path=args.output_hub_path,
+ revision=args.revision,
+ )
+
+
+if __name__ == "__main__":
+ main()
diff --git a/third_party/transformers/src/transformers/models/eomt/image_processing_eomt.py b/third_party/transformers/src/transformers/models/eomt/image_processing_eomt.py
new file mode 100644
index 0000000000000000000000000000000000000000..7d8099fa24879be51f7c760f48a92ec500ac8b34
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/eomt/image_processing_eomt.py
@@ -0,0 +1,670 @@
+# Copyright 2025 Mobile Perception Systems Lab at TU/e and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Image processor class for EoMT."""
+
+import math
+from typing import Union
+
+import numpy as np
+import torch
+from torchvision.transforms.v2 import functional as tvF
+
+from ...image_processing_backends import TorchvisionBackend
+from ...image_processing_utils import BatchFeature
+from ...image_transforms import get_size_with_aspect_ratio, group_images_by_shape, reorder_images
+from ...image_utils import (
+ IMAGENET_DEFAULT_MEAN,
+ IMAGENET_DEFAULT_STD,
+ ChannelDimension,
+ ImageInput,
+ PILImageResampling,
+ SizeDict,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import (
+ TensorType,
+ auto_docstring,
+ filter_out_non_signature_kwargs,
+)
+
+
+class EomtImageProcessorKwargs(ImagesKwargs, total=False):
+ r"""
+ do_split_image (`bool`, *optional*, defaults to `self.do_split_image`):
+ Whether to split the input images into overlapping patches for semantic segmentation. If set to `True`, the
+ input images will be split into patches of size `size["shortest_edge"]` with an overlap between patches.
+ Otherwise, the input images will be padded to the target size.
+ ignore_index (`int`, *optional*, defaults to `self.ignore_index`):
+ Label to be assigned to background pixels in segmentation maps. If provided, segmentation map pixels
+ denoted with 0 (background) will be replaced with `ignore_index`.
+ """
+
+ do_split_image: bool
+ ignore_index: int | None
+
+
+# Adapted from transformers.models.maskformer.image_processing_maskformer_fast.convert_segmentation_map_to_binary_masks_fast
+def convert_segmentation_map_to_binary_masks_fast(
+ segmentation_map: "torch.Tensor",
+ instance_id_to_semantic_id: dict[int, int] | None = None,
+ ignore_index: int | None = None,
+):
+ if ignore_index is not None:
+ segmentation_map = torch.where(segmentation_map == 0, ignore_index, segmentation_map - 1)
+
+ all_labels = torch.unique(segmentation_map)
+
+ if ignore_index is not None:
+ all_labels = all_labels[all_labels != ignore_index] # drop background label if applicable
+
+ binary_masks = [(segmentation_map == i) for i in all_labels]
+ if binary_masks:
+ binary_masks = torch.stack(binary_masks, dim=0)
+ else:
+ binary_masks = torch.zeros((0, *segmentation_map.shape), device=segmentation_map.device)
+
+ # Convert instance ids to class ids
+ if instance_id_to_semantic_id is not None:
+ labels = torch.zeros(all_labels.shape[0], device=segmentation_map.device)
+
+ for i, label in enumerate(all_labels):
+ class_id = instance_id_to_semantic_id[(label.item() + 1 if ignore_index is not None else label.item())]
+ labels[i] = class_id - 1 if ignore_index is not None else class_id
+ else:
+ labels = all_labels
+ return binary_masks.float(), labels.long()
+
+
+def get_target_size(size_dict: dict[str, int]) -> tuple[int, int]:
+ """Returns the height and width from a size dict."""
+ target_height = size_dict["shortest_edge"]
+ target_width = size_dict["longest_edge"] or target_height
+
+ return target_height, target_width
+
+
+def reorder_patches_and_offsets(
+ patches: list[torch.Tensor], offsets: list[list[int]]
+) -> tuple[list[torch.Tensor], list[list[int]]]:
+ """Sorts patches and offsets according to the original image index."""
+
+ combined = list(zip(offsets, patches))
+ combined.sort(key=lambda x: x[0][0])
+ sorted_offsets, sorted_patches = zip(*combined)
+
+ return list(sorted_patches), list(sorted_offsets)
+
+
+def remove_low_and_no_objects(masks, scores, labels, object_mask_threshold, num_labels):
+ """
+ Binarize the given masks using `object_mask_threshold`, it returns the associated values of `masks`, `scores` and
+ `labels`.
+
+ Args:
+ masks (`torch.Tensor`):
+ A tensor of shape `(num_queries, height, width)`.
+ scores (`torch.Tensor`):
+ A tensor of shape `(num_queries)`.
+ labels (`torch.Tensor`):
+ A tensor of shape `(num_queries)`.
+ object_mask_threshold (`float`):
+ A number between 0 and 1 used to binarize the masks.
+ Raises:
+ `ValueError`: Raised when the first dimension doesn't match in all input tensors.
+ Returns:
+ `tuple[`torch.Tensor`, `torch.Tensor`, `torch.Tensor`]`: The `masks`, `scores` and `labels` without the region
+ < `object_mask_threshold`.
+ """
+ if not (masks.shape[0] == scores.shape[0] == labels.shape[0]):
+ raise ValueError("mask, scores and labels must have the same shape!")
+
+ to_keep = labels.ne(num_labels) & (scores > object_mask_threshold)
+
+ return masks[to_keep], scores[to_keep], labels[to_keep]
+
+
+def check_segment_validity(mask_labels, mask_probs, k, mask_threshold=0.5, overlap_mask_area_threshold=0.8):
+ # Get the mask associated with the k class
+ mask_k = mask_labels == k
+ mask_k_area = mask_k.sum()
+
+ # Compute the area of all the stuff in query k
+ original_mask = mask_probs[k] >= mask_threshold
+ original_area = original_mask.sum()
+
+ final_mask = mask_k & original_mask
+ final_mask_area = final_mask.sum()
+
+ mask_exists = mask_k_area > 0 and original_area > 0 and final_mask_area > 0
+
+ if mask_exists:
+ area_ratio = mask_k_area / original_area
+ if not area_ratio.item() > overlap_mask_area_threshold:
+ mask_exists = False
+
+ return mask_exists, final_mask
+
+
+def compute_segments(
+ mask_probs,
+ pred_scores,
+ pred_labels,
+ stuff_classes,
+ mask_threshold: float = 0.5,
+ overlap_mask_area_threshold: float = 0.8,
+ target_size: tuple[int, int] | None = None,
+):
+ height = mask_probs.shape[1] if target_size is None else target_size[0]
+ width = mask_probs.shape[2] if target_size is None else target_size[1]
+
+ segmentation = torch.zeros((height, width), dtype=torch.long, device=mask_probs.device) - 1
+ segments: list[dict] = []
+
+ # Compute per-pixel assignment based on weighted mask scores
+ mask_probs = mask_probs.sigmoid()
+ mask_labels = (pred_scores[:, None, None] * mask_probs).argmax(0)
+
+ # Keep track of instances of each class
+ current_segment_id = 0
+ stuff_memory_list: dict[str, int] = {}
+
+ for k in range(pred_labels.shape[0]):
+ pred_class = pred_labels[k].item()
+
+ # Check if mask exists and large enough to be a segment
+ mask_exists, final_mask = check_segment_validity(
+ mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold
+ )
+
+ if not mask_exists:
+ continue
+
+ if stuff_classes and pred_class in stuff_classes:
+ if pred_class in stuff_memory_list:
+ segmentation[final_mask] = stuff_memory_list[pred_class]
+ continue
+ else:
+ stuff_memory_list[pred_class] = current_segment_id
+
+ segmentation[final_mask] = current_segment_id
+ segment_score = round(pred_scores[k].item(), 6)
+ segments.append(
+ {
+ "id": current_segment_id,
+ "label_id": pred_class,
+ "score": segment_score,
+ }
+ )
+ current_segment_id += 1
+ return segmentation, segments
+
+
+@auto_docstring
+class EomtImageProcessor(TorchvisionBackend):
+ valid_kwargs = EomtImageProcessorKwargs
+ resample = PILImageResampling.BILINEAR
+ image_mean = IMAGENET_DEFAULT_MEAN
+ image_std = IMAGENET_DEFAULT_STD
+ size = {"shortest_edge": 640, "longest_edge": 640}
+ default_to_square = False
+ do_resize = True
+ do_rescale = True
+ do_normalize = True
+ do_split_image = False
+ do_pad = False
+ ignore_index = None
+
+ def __init__(self, **kwargs: Unpack[EomtImageProcessorKwargs]):
+ super().__init__(**kwargs)
+
+ def _split_image(self, images: torch.Tensor, size: SizeDict, image_indices: list[int]) -> tuple[list, list]:
+ """Slices an image into overlapping patches for semantic segmentation."""
+
+ patches, patch_offsets = [], []
+
+ _, _, height, width = images.shape
+ patch_size = size.shortest_edge
+
+ longer_side = max(height, width)
+ num_patches = math.ceil(longer_side / patch_size)
+ total_overlap = num_patches * patch_size - longer_side
+ overlap_per_patch = total_overlap / (num_patches - 1) if num_patches > 1 else 0
+
+ for i in range(num_patches):
+ start = int(i * (patch_size - overlap_per_patch))
+ end = start + patch_size
+
+ if height > width:
+ batch_patch = images[:, :, start:end, :]
+ else:
+ batch_patch = images[:, :, :, start:end]
+
+ for batch_idx, single in enumerate(torch.unbind(batch_patch, dim=0)):
+ patches.append(single)
+ patch_offsets.append([image_indices[batch_idx], start, end])
+
+ return patches, patch_offsets
+
+ def _pad(self, images: torch.Tensor, size: SizeDict) -> torch.Tensor:
+ """Pads the image to the target size using zero padding."""
+ _, _, height, width = images.shape
+
+ target_height, target_width = get_target_size(
+ {"shortest_edge": size.shortest_edge, "longest_edge": size.longest_edge or size.shortest_edge}
+ )
+ pad_h = max(0, target_height - height)
+ pad_w = max(0, target_width - width)
+ padding = (0, pad_w, 0, pad_h)
+
+ padded_images = torch.nn.functional.pad(images, padding, mode="constant", value=0.0)
+ return padded_images
+
+ @auto_docstring
+ def preprocess(
+ self,
+ images: ImageInput,
+ segmentation_maps: list[torch.Tensor] | None = None,
+ instance_id_to_semantic_id: dict[int, int] | None = None,
+ **kwargs: Unpack[EomtImageProcessorKwargs],
+ ) -> BatchFeature:
+ r"""
+ segmentation_maps (`ImageInput`, *optional*):
+ The segmentation maps to preprocess for corresponding images.
+ instance_id_to_semantic_id (`list[dict[int, int]]` or `dict[int, int]`, *optional*):
+ A mapping between object instance ids and class ids.
+ """
+ return super().preprocess(images, segmentation_maps, instance_id_to_semantic_id, **kwargs)
+
+ def _preprocess_image_like_inputs(
+ self,
+ images: ImageInput,
+ segmentation_maps: ImageInput | None,
+ instance_id_to_semantic_id: dict[int, int] | None,
+ do_convert_rgb: bool,
+ input_data_format: ChannelDimension,
+ return_tensors: str | TensorType | None,
+ device: Union[str, "torch.device"] | None = None,
+ **kwargs: Unpack[EomtImageProcessorKwargs],
+ ) -> BatchFeature:
+ """
+ Preprocess image-like inputs.
+ """
+ images = self._prepare_image_like_inputs(
+ images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format, device=device
+ )
+ ignore_index = kwargs.pop("ignore_index", None)
+ images_kwargs = kwargs.copy()
+ data = {}
+ processed_images, patch_offsets = self._preprocess(images, **images_kwargs)
+ data["pixel_values"] = processed_images
+ data["patch_offsets"] = patch_offsets
+
+ if segmentation_maps is not None:
+ processed_segmentation_maps = self._prepare_image_like_inputs(
+ images=segmentation_maps,
+ expected_ndims=2,
+ do_convert_rgb=False,
+ input_data_format=ChannelDimension.FIRST,
+ )
+
+ segmentation_maps_kwargs = kwargs.copy()
+ segmentation_maps_kwargs.update(
+ {
+ "do_normalize": False,
+ "do_rescale": False,
+ # Nearest interpolation is used for segmentation maps instead of BILINEAR.
+ "resample": PILImageResampling.NEAREST,
+ }
+ )
+
+ processed_segmentation_maps, _ = self._preprocess(
+ images=processed_segmentation_maps, **segmentation_maps_kwargs
+ )
+ processed_segmentation_maps = [
+ segmentation_map.squeeze(0).to(torch.int64) for segmentation_map in processed_segmentation_maps
+ ]
+ # Convert to list of binary masks and labels
+ mask_labels, class_labels = [], []
+ for idx, segmentation_map in enumerate(processed_segmentation_maps):
+ if isinstance(instance_id_to_semantic_id, list):
+ instance_id = instance_id_to_semantic_id[idx]
+ else:
+ instance_id = instance_id_to_semantic_id
+ # Use instance2class_id mapping per image
+ masks, classes = convert_segmentation_map_to_binary_masks_fast(
+ segmentation_map,
+ instance_id,
+ ignore_index=ignore_index,
+ )
+
+ mask_labels.append(masks)
+ class_labels.append(classes)
+
+ # we cannot batch them since they don't share a common class size
+ data["mask_labels"] = mask_labels
+ data["class_labels"] = class_labels
+
+ if patch_offsets:
+ data["patch_offsets"] = [torch.tensor(offsets) for offsets in patch_offsets]
+
+ return BatchFeature(
+ data=data,
+ tensor_type=return_tensors,
+ skip_tensor_conversion=["patch_offsets", "mask_labels", "class_labels"],
+ )
+
+ def _preprocess(
+ self,
+ images: list["torch.Tensor"],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ do_split_image: bool,
+ do_pad: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ disable_grouping: bool | None,
+ **kwargs,
+ ):
+ """Preprocesses the input images and masks if provided."""
+ patch_offsets = []
+
+ grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+ resized_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ if do_resize:
+ stacked_images = self.resize(image=stacked_images, size=size, resample=resample)
+ resized_images_grouped[shape] = stacked_images
+ images = reorder_images(resized_images_grouped, grouped_images_index)
+
+ if do_split_image:
+ grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+ patches, patch_offsets = [], []
+ for shape, stacked_images in grouped_images.items():
+ original_indices = [
+ original_idx for original_idx, (img_shape, _) in grouped_images_index.items() if img_shape == shape
+ ]
+ split_patches, offsets = self._split_image(stacked_images, size, original_indices)
+ patches.extend(split_patches)
+ patch_offsets.extend(offsets)
+ images, patch_offsets = reorder_patches_and_offsets(patches, patch_offsets)
+
+ if do_pad:
+ grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+ padded_grouped = {
+ shape: self._pad(stacked_images, size) for shape, stacked_images in grouped_images.items()
+ }
+ images = reorder_images(padded_grouped, grouped_images_index)
+
+ grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+ processed_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ stacked_images = self.rescale_and_normalize(
+ stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
+ )
+ processed_images_grouped[shape] = stacked_images
+ processed_images = reorder_images(processed_images_grouped, grouped_images_index)
+
+ return processed_images, patch_offsets
+
+ def merge_image_patches(
+ self,
+ segmentation_logits: torch.Tensor,
+ patch_offsets: list[tuple[int, int, int]],
+ target_sizes: list[tuple[int, int]],
+ size: dict[str, int],
+ ) -> list[torch.Tensor]:
+ """
+ Reconstructs full-size semantic segmentation logits from patch predictions.
+
+ Args:
+ segmentation_logits (`torch.Tensor`):
+ A tensor of shape `(num_patches, num_classes, patch_height, patch_width)` representing predicted logits
+ for each image patch.
+ patch_offsets (`list[tuple[int, int, int]]`):
+ A list of tuples where each tuple contains:
+ - `image_index` (int): Index of the original image this patch belongs to.
+ - `start` (int): Start pixel index of the patch along the long dimension (height or width).
+ - `end` (int): End pixel index of the patch along the long dimension.
+ target_sizes (`list[tuple[int, int]]`):
+ list of original (height, width) dimensions for each image before preprocessing.
+ size (`dict[str, int]`):
+ A size dict which was used to resize.
+ """
+ num_classes = segmentation_logits.shape[1]
+ aggregated_logits = []
+ patch_counts = []
+
+ for image_size in target_sizes:
+ height, width = get_size_with_aspect_ratio(image_size, size["shortest_edge"], size["longest_edge"])
+ aggregated_logits.append(torch.zeros((num_classes, height, width), device=segmentation_logits.device))
+ patch_counts.append(torch.zeros((num_classes, height, width), device=segmentation_logits.device))
+
+ # Stitch patches back into full-sized logit maps
+ for patch_idx, (image_idx, patch_start, patch_end) in enumerate(patch_offsets):
+ if target_sizes[image_idx][0] > target_sizes[image_idx][1]:
+ aggregated_logits[image_idx][:, patch_start:patch_end, :] += segmentation_logits[patch_idx]
+ patch_counts[image_idx][:, patch_start:patch_end, :] += 1
+ else:
+ aggregated_logits[image_idx][:, :, patch_start:patch_end] += segmentation_logits[patch_idx]
+ patch_counts[image_idx][:, :, patch_start:patch_end] += 1
+
+ # Normalize and resize logits to original image size
+ reconstructed_logits = []
+ for idx, (logit_sum, count) in enumerate(zip(aggregated_logits, patch_counts)):
+ averaged_logits = logit_sum / count.clamp(min=1)
+ resized_logits = torch.nn.functional.interpolate(
+ averaged_logits[None, ...],
+ size=target_sizes[idx],
+ mode="bilinear",
+ align_corners=False,
+ )[0]
+
+ reconstructed_logits.append(resized_logits)
+
+ return reconstructed_logits
+
+ def unpad_image(
+ self,
+ segmentation_logits: torch.Tensor,
+ target_sizes: list[tuple[int, int]],
+ size: dict[str, int],
+ ) -> list[torch.Tensor]:
+ """Restores panoptic segmentation logits to their original image resolutions."""
+
+ resized_logits = []
+
+ for idx, original_size in enumerate(target_sizes):
+ target_height, target_width = get_size_with_aspect_ratio(
+ original_size, size["shortest_edge"], size["longest_edge"]
+ )
+ cropped_logits = segmentation_logits[idx][:, :target_height, :target_width]
+ upsampled_logits = torch.nn.functional.interpolate(
+ cropped_logits[None, ...], size=original_size, mode="bilinear", align_corners=False
+ )[0]
+ resized_logits.append(upsampled_logits)
+ return resized_logits
+
+ def post_process_semantic_segmentation(
+ self,
+ outputs,
+ target_sizes: list[tuple[int, int]],
+ size: dict[str, int] | None = None,
+ ) -> np.ndarray:
+ """Post-processes model outputs into final semantic segmentation prediction."""
+
+ size = size if size is not None else self.size
+
+ masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width]
+ class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1]
+ patch_offsets = outputs.patch_offsets
+
+ output_size = get_target_size(size)
+ masks_queries_logits = torch.nn.functional.interpolate(
+ masks_queries_logits,
+ size=output_size,
+ mode="bilinear",
+ )
+
+ # Remove the null class `[..., :-1]`
+ masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
+ masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width]
+
+ segmentation_logits = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)
+
+ if patch_offsets:
+ output_logits = self.merge_image_patches(segmentation_logits, patch_offsets, target_sizes, size)
+ else:
+ output_logits = []
+
+ for idx in range(len(segmentation_logits)):
+ resized_logits = torch.nn.functional.interpolate(
+ segmentation_logits[idx].unsqueeze(dim=0),
+ size=target_sizes[idx],
+ mode="bilinear",
+ align_corners=False,
+ )
+ output_logits.append(resized_logits[0])
+
+ preds = [logit.argmax(dim=0) for logit in output_logits]
+ return preds
+
+ def post_process_panoptic_segmentation(
+ self,
+ outputs,
+ target_sizes: list[tuple[int, int]],
+ threshold: float = 0.8,
+ mask_threshold: float = 0.5,
+ overlap_mask_area_threshold: float = 0.8,
+ stuff_classes: list[int] | None = None,
+ size: dict[str, int] | None = None,
+ ):
+ """Post-processes model outputs into final panoptic segmentation prediction."""
+
+ size = size if size is not None else self.size
+
+ masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width]
+ class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1]
+
+ batch_size = class_queries_logits.shape[0]
+ num_labels = class_queries_logits.shape[-1] - 1
+
+ output_size = get_target_size(size)
+ masks_queries_logits = torch.nn.functional.interpolate(
+ masks_queries_logits,
+ size=output_size,
+ mode="bilinear",
+ )
+
+ mask_probs_batch = self.unpad_image(masks_queries_logits, target_sizes, size)
+ pred_scores_batch, pred_labels_batch = class_queries_logits.softmax(dim=-1).max(-1)
+
+ results: list = []
+
+ for i in range(batch_size):
+ mask_probs, pred_scores, pred_labels = remove_low_and_no_objects(
+ mask_probs_batch[i], pred_scores_batch[i], pred_labels_batch[i], threshold, num_labels
+ )
+
+ # No mask found
+ if mask_probs.shape[0] <= 0:
+ height, width = target_sizes[i] if target_sizes is not None else mask_probs.shape[1:]
+ segmentation = torch.zeros((height, width)) - 1
+ results.append({"segmentation": segmentation, "segments_info": []})
+ continue
+
+ segmentation, segments = compute_segments(
+ mask_probs=mask_probs,
+ pred_scores=pred_scores,
+ pred_labels=pred_labels,
+ stuff_classes=stuff_classes,
+ mask_threshold=mask_threshold,
+ overlap_mask_area_threshold=overlap_mask_area_threshold,
+ target_size=target_sizes[i] if target_sizes is not None else None,
+ )
+
+ results.append({"segmentation": segmentation, "segments_info": segments})
+ return results
+
+ @filter_out_non_signature_kwargs()
+ def post_process_instance_segmentation(
+ self,
+ outputs,
+ target_sizes: list[tuple[int, int]],
+ threshold: float = 0.8,
+ size: dict[str, int] | None = None,
+ ):
+ """Post-processes model outputs into Instance Segmentation Predictions."""
+
+ size = size if size is not None else self.size
+
+ masks_queries_logits = outputs.masks_queries_logits
+ class_queries_logits = outputs.class_queries_logits
+
+ output_size = get_target_size(size)
+ masks_queries_logits = torch.nn.functional.interpolate(
+ masks_queries_logits,
+ size=output_size,
+ mode="bilinear",
+ )
+
+ mask_probs_batch = self.unpad_image(masks_queries_logits, target_sizes, size)
+
+ device = masks_queries_logits.device
+ batch_size = class_queries_logits.shape[0]
+ num_queries = class_queries_logits.shape[-2]
+
+ results = []
+
+ for i in range(batch_size):
+ mask_pred = mask_probs_batch[i]
+ mask_class = class_queries_logits[i]
+
+ # Remove the null class `[..., :-1]`
+ scores, pred_classes = mask_class.softmax(dim=-1)[..., :-1].max(-1)
+ pred_masks = (mask_pred > 0).float()
+
+ # Calculate average mask prob
+ mask_scores = (mask_pred.sigmoid().flatten(1) * pred_masks.flatten(1)).sum(1) / (
+ pred_masks.flatten(1).sum(1) + 1e-6
+ )
+ pred_scores = scores * mask_scores
+
+ segmentation = torch.zeros(target_sizes[i], device=device) - 1
+
+ instance_maps, segments = [], []
+ current_segment_id = 0
+ for j in range(num_queries):
+ score = pred_scores[j].item()
+
+ if not torch.all(pred_masks[j] == 0) and score >= threshold:
+ segmentation[pred_masks[j] == 1] = current_segment_id
+ segments.append(
+ {
+ "id": current_segment_id,
+ "label_id": pred_classes[j].item(),
+ "score": round(score, 6),
+ }
+ )
+ current_segment_id += 1
+ instance_maps.append(pred_masks[j])
+
+ results.append({"segmentation": segmentation, "segments_info": segments})
+ return results
+
+
+__all__ = ["EomtImageProcessor"]
diff --git a/third_party/transformers/src/transformers/models/eomt/image_processing_pil_eomt.py b/third_party/transformers/src/transformers/models/eomt/image_processing_pil_eomt.py
new file mode 100644
index 0000000000000000000000000000000000000000..841537838a09284eacd3df8a5857f5877743b8c0
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/eomt/image_processing_pil_eomt.py
@@ -0,0 +1,642 @@
+# Copyright 2025 Mobile Perception Systems Lab at TU/e and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Image processor class for EoMT."""
+
+import math
+
+import numpy as np
+
+from ...image_processing_backends import PilBackend
+from ...image_processing_utils import BatchFeature
+from ...image_transforms import PaddingMode, get_size_with_aspect_ratio
+from ...image_transforms import pad as np_pad
+from ...image_utils import (
+ IMAGENET_DEFAULT_MEAN,
+ IMAGENET_DEFAULT_STD,
+ ChannelDimension,
+ ImageInput,
+ PILImageResampling,
+ SizeDict,
+ get_image_size,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring, is_torch_available
+from ...utils.import_utils import requires
+
+
+if is_torch_available():
+ import torch
+
+
+# Adapted from transformers.models.maskformer.image_processing_maskformer.convert_segmentation_map_to_binary_masks
+def convert_segmentation_map_to_binary_masks(
+ segmentation_map: np.ndarray,
+ instance_id_to_semantic_id: dict[int, int] | None = None,
+ ignore_index: int | None = None,
+):
+ if ignore_index is not None:
+ segmentation_map = np.where(segmentation_map == 0, ignore_index, segmentation_map - 1)
+
+ # Get unique ids (class or instance ids based on input)
+ all_labels = np.unique(segmentation_map)
+
+ # Drop background label if applicable
+ if ignore_index is not None:
+ all_labels = all_labels[all_labels != ignore_index]
+
+ # Generate a binary mask for each object instance
+ binary_masks = [(segmentation_map == i) for i in all_labels]
+
+ # Stack the binary masks
+ if binary_masks:
+ binary_masks = np.stack(binary_masks, axis=0)
+ else:
+ binary_masks = np.zeros((0, *segmentation_map.shape))
+
+ # Convert instance ids to class ids
+ if instance_id_to_semantic_id is not None:
+ labels = np.zeros(all_labels.shape[0])
+
+ for label in all_labels:
+ class_id = instance_id_to_semantic_id[label + 1 if ignore_index is not None else label]
+ labels[all_labels == label] = class_id - 1 if ignore_index is not None else class_id
+ else:
+ labels = all_labels
+
+ return binary_masks.astype(np.float32), labels.astype(np.int64)
+
+
+# Adapted from transformers.models.eomt.image_processing_eomt.check_segment_validity
+def check_segment_validity(mask_labels, mask_probs, k, mask_threshold=0.5, overlap_mask_area_threshold=0.8):
+ # Get the mask associated with the k class
+ mask_k = mask_labels == k
+ mask_k_area = mask_k.sum()
+
+ # Compute the area of all the stuff in query k
+ original_mask = mask_probs[k] >= mask_threshold
+ original_area = original_mask.sum()
+
+ final_mask = mask_k & original_mask
+ final_mask_area = final_mask.sum()
+
+ mask_exists = mask_k_area > 0 and original_area > 0 and final_mask_area > 0
+
+ if mask_exists:
+ area_ratio = mask_k_area / original_area
+ if not area_ratio.item() > overlap_mask_area_threshold:
+ mask_exists = False
+
+ return mask_exists, final_mask
+
+
+# Adapted from transformers.models.eomt.image_processing_eomt.EomtImageProcessorKwargs
+class EomtImageProcessorKwargs(ImagesKwargs, total=False):
+ r"""
+ do_split_image (`bool`, *optional*, defaults to `self.do_split_image`):
+ Whether to split the input images into overlapping patches for semantic segmentation. If set to `True`, the
+ input images will be split into patches of size `size["shortest_edge"]` with an overlap between patches.
+ Otherwise, the input images will be padded to the target size.
+ ignore_index (`int`, *optional*, defaults to `self.ignore_index`):
+ Label to be assigned to background pixels in segmentation maps. If provided, segmentation map pixels
+ denoted with 0 (background) will be replaced with `ignore_index`.
+ """
+
+ do_split_image: bool
+ ignore_index: int | None
+
+
+# Adapted from transformers.models.eomt.image_processing_eomt.compute_segments
+def compute_segments(
+ mask_probs,
+ pred_scores,
+ pred_labels,
+ stuff_classes,
+ mask_threshold: float = 0.5,
+ overlap_mask_area_threshold: float = 0.8,
+ target_size: tuple[int, int] | None = None,
+):
+ height = mask_probs.shape[1] if target_size is None else target_size[0]
+ width = mask_probs.shape[2] if target_size is None else target_size[1]
+
+ segmentation = torch.zeros((height, width), dtype=torch.long, device=mask_probs.device) - 1
+ segments: list[dict] = []
+
+ # Compute per-pixel assignment based on weighted mask scores
+ mask_probs = mask_probs.sigmoid()
+ mask_labels = (pred_scores[:, None, None] * mask_probs).argmax(0)
+
+ # Keep track of instances of each class
+ current_segment_id = 0
+ stuff_memory_list: dict[str, int] = {}
+
+ for k in range(pred_labels.shape[0]):
+ pred_class = pred_labels[k].item()
+
+ # Check if mask exists and large enough to be a segment
+ mask_exists, final_mask = check_segment_validity(
+ mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold
+ )
+
+ if not mask_exists:
+ continue
+
+ if stuff_classes and pred_class in stuff_classes:
+ if pred_class in stuff_memory_list:
+ segmentation[final_mask] = stuff_memory_list[pred_class]
+ continue
+ else:
+ stuff_memory_list[pred_class] = current_segment_id
+
+ segmentation[final_mask] = current_segment_id
+ segment_score = round(pred_scores[k].item(), 6)
+ segments.append(
+ {
+ "id": current_segment_id,
+ "label_id": pred_class,
+ "score": segment_score,
+ }
+ )
+ current_segment_id += 1
+ return segmentation, segments
+
+
+# Adapted from transformers.models.eomt.image_processing_eomt.get_target_size
+def get_target_size(size_dict: dict[str, int]) -> tuple[int, int]:
+ """Returns the height and width from a size dict."""
+ target_height = size_dict["shortest_edge"]
+ target_width = size_dict["longest_edge"] or target_height
+
+ return target_height, target_width
+
+
+# Adapted from transformers.models.eomt.image_processing_eomt.remove_low_and_no_objects
+def remove_low_and_no_objects(masks, scores, labels, object_mask_threshold, num_labels):
+ """
+ Binarize the given masks using `object_mask_threshold`, it returns the associated values of `masks`, `scores` and
+ `labels`.
+
+ Args:
+ masks (`torch.Tensor`):
+ A tensor of shape `(num_queries, height, width)`.
+ scores (`torch.Tensor`):
+ A tensor of shape `(num_queries)`.
+ labels (`torch.Tensor`):
+ A tensor of shape `(num_queries)`.
+ object_mask_threshold (`float`):
+ A number between 0 and 1 used to binarize the masks.
+ Raises:
+ `ValueError`: Raised when the first dimension doesn't match in all input tensors.
+ Returns:
+ `tuple[`torch.Tensor`, `torch.Tensor`, `torch.Tensor`]`: The `masks`, `scores` and `labels` without the region
+ < `object_mask_threshold`.
+ """
+ if not (masks.shape[0] == scores.shape[0] == labels.shape[0]):
+ raise ValueError("mask, scores and labels must have the same shape!")
+
+ to_keep = labels.ne(num_labels) & (scores > object_mask_threshold)
+
+ return masks[to_keep], scores[to_keep], labels[to_keep]
+
+
+@auto_docstring
+@requires(backends=("torch",))
+class EomtImageProcessorPil(PilBackend):
+ valid_kwargs = EomtImageProcessorKwargs
+ resample = PILImageResampling.BILINEAR
+ image_mean = IMAGENET_DEFAULT_MEAN
+ image_std = IMAGENET_DEFAULT_STD
+ size = {"shortest_edge": 640, "longest_edge": 640}
+ default_to_square = False
+ do_resize = True
+ do_rescale = True
+ do_normalize = True
+ do_split_image = False
+ do_pad = False
+ ignore_index = None
+
+ def __init__(self, **kwargs: Unpack[EomtImageProcessorKwargs]):
+ super().__init__(**kwargs)
+
+ def _split_image(self, image: np.ndarray, size: SizeDict, image_index: int) -> tuple[list, list]:
+ """Slices an image into overlapping patches for semantic segmentation."""
+
+ patches, patch_offsets = [], []
+
+ height, width = get_image_size(image, channel_dim=ChannelDimension.FIRST)
+ patch_size = size.shortest_edge
+
+ longer_side = max(height, width)
+ num_patches = math.ceil(longer_side / patch_size)
+ total_overlap = num_patches * patch_size - longer_side
+ overlap_per_patch = total_overlap / (num_patches - 1) if num_patches > 1 else 0
+
+ for i in range(num_patches):
+ start = int(i * (patch_size - overlap_per_patch))
+ end = start + patch_size
+
+ if height > width:
+ patch = image[:, start:end, :]
+ else:
+ patch = image[:, :, start:end]
+
+ patches.append(patch)
+ patch_offsets.append([image_index, start, end])
+
+ return patches, patch_offsets
+
+ def _pad(self, image: np.ndarray, size: SizeDict) -> np.ndarray:
+ """Pads the image to the target size using zero padding."""
+
+ height, width = get_image_size(image, channel_dim=ChannelDimension.FIRST)
+
+ target_height, target_width = get_target_size(
+ {"shortest_edge": size.shortest_edge, "longest_edge": size.longest_edge or size.shortest_edge}
+ )
+ pad_h = max(0, target_height - height)
+ pad_w = max(0, target_width - width)
+
+ padding = ((0, pad_h), (0, pad_w))
+ padded_image = np_pad(
+ image,
+ padding,
+ mode=PaddingMode.CONSTANT,
+ constant_values=0.0,
+ data_format=ChannelDimension.FIRST,
+ input_data_format=ChannelDimension.FIRST,
+ )
+ return padded_image
+
+ @auto_docstring
+ def preprocess(
+ self,
+ images: ImageInput,
+ segmentation_maps: "list[torch.Tensor] | None" = None,
+ instance_id_to_semantic_id: dict[int, int] | None = None,
+ **kwargs: Unpack[EomtImageProcessorKwargs],
+ ) -> BatchFeature:
+ r"""
+ segmentation_maps (`ImageInput`, *optional*):
+ The segmentation maps to preprocess for corresponding images.
+ instance_id_to_semantic_id (`list[dict[int, int]]` or `dict[int, int]`, *optional*):
+ A mapping between object instance ids and class ids.
+ """
+ return super().preprocess(images, segmentation_maps, instance_id_to_semantic_id, **kwargs)
+
+ def _preprocess_image_like_inputs(
+ self,
+ images: ImageInput,
+ segmentation_maps: ImageInput | None,
+ instance_id_to_semantic_id: dict[int, int] | None,
+ do_convert_rgb: bool,
+ input_data_format: ChannelDimension,
+ return_tensors: str | TensorType | None,
+ **kwargs: Unpack[EomtImageProcessorKwargs],
+ ) -> BatchFeature:
+ """
+ Preprocess image-like inputs.
+ """
+ images = self._prepare_image_like_inputs(
+ images=images, do_convert_rgb=do_convert_rgb, input_data_format=input_data_format
+ )
+ ignore_index = kwargs.pop("ignore_index", None)
+ images_kwargs = kwargs.copy()
+ processed_images, patch_offsets = self._preprocess(images, **images_kwargs)
+ data = {}
+ data["pixel_values"] = processed_images
+ data["patch_offsets"] = patch_offsets
+
+ if segmentation_maps is not None:
+ processed_segmentation_maps = self._prepare_image_like_inputs(
+ images=segmentation_maps,
+ expected_ndims=2,
+ do_convert_rgb=False,
+ input_data_format=ChannelDimension.FIRST,
+ )
+
+ segmentation_maps_kwargs = kwargs.copy()
+ segmentation_maps_kwargs.update(
+ {
+ "do_normalize": False,
+ "do_rescale": False,
+ # Nearest interpolation is used for segmentation maps instead of BILINEAR.
+ "resample": PILImageResampling.NEAREST,
+ }
+ )
+
+ processed_segmentation_maps, _ = self._preprocess(
+ images=processed_segmentation_maps, **segmentation_maps_kwargs
+ )
+ processed_segmentation_maps = [
+ segmentation_map.squeeze(0).astype(np.int64) for segmentation_map in processed_segmentation_maps
+ ]
+
+ # Convert to list of binary masks and labels
+ mask_labels, class_labels = [], []
+ for idx, segmentation_map in enumerate(processed_segmentation_maps):
+ if isinstance(instance_id_to_semantic_id, list):
+ instance_id = instance_id_to_semantic_id[idx]
+ else:
+ instance_id = instance_id_to_semantic_id
+ # Use instance2class_id mapping per image
+ masks, classes = convert_segmentation_map_to_binary_masks(
+ segmentation_map, instance_id, ignore_index=ignore_index
+ )
+
+ mask_labels.append(torch.from_numpy(masks))
+ class_labels.append(torch.from_numpy(classes))
+
+ # we cannot batch them since they don't share a common class size
+ data["mask_labels"] = mask_labels
+ data["class_labels"] = class_labels
+
+ if patch_offsets:
+ data["patch_offsets"] = [torch.tensor(offsets) for offsets in patch_offsets]
+
+ return BatchFeature(
+ data=data,
+ tensor_type=return_tensors,
+ skip_tensor_conversion=["patch_offsets", "mask_labels", "class_labels"],
+ )
+
+ def _preprocess(
+ self,
+ images: list[np.ndarray],
+ do_resize: bool,
+ size: SizeDict,
+ resample: PILImageResampling | None,
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ do_split_image: bool,
+ do_pad: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ **kwargs,
+ ):
+ """Preprocesses the input images and masks if provided."""
+ processed_images, patch_offsets = [], []
+
+ # Resize images
+ resized_images = []
+ for image in images:
+ if do_resize:
+ image = self.resize(image, size, resample)
+ resized_images.append(image)
+
+ # Split images into patches if requested
+ if do_split_image:
+ for idx, img in enumerate(resized_images):
+ patches, offsets = self._split_image(img, size, idx)
+ processed_images.extend(patches)
+ patch_offsets.extend(offsets)
+ images = processed_images
+ else:
+ images = resized_images
+
+ # Pad images if requested
+ if do_pad:
+ images = [self._pad(img, size) for img in images]
+
+ # Rescale and normalize
+ processed_images = []
+ for image in images:
+ if do_rescale:
+ image = self.rescale(image, rescale_factor)
+ if do_normalize:
+ image = self.normalize(image, image_mean, image_std)
+ processed_images.append(image)
+
+ return processed_images, patch_offsets
+
+ def merge_image_patches(
+ self,
+ segmentation_logits: "torch.Tensor",
+ patch_offsets: list[tuple[int, int, int]],
+ target_sizes: list[tuple[int, int]],
+ size: dict[str, int],
+ ) -> "list[torch.Tensor]":
+ """
+ Reconstructs full-size semantic segmentation logits from patch predictions.
+
+ Args:
+ segmentation_logits (`torch.Tensor`):
+ A tensor of shape `(num_patches, num_classes, patch_height, patch_width)` representing predicted logits
+ for each image patch.
+ patch_offsets (`list[tuple[int, int, int]]`):
+ A list of tuples where each tuple contains:
+ - `image_index` (int): Index of the original image this patch belongs to.
+ - `start` (int): Start pixel index of the patch along the long dimension (height or width).
+ - `end` (int): End pixel index of the patch along the long dimension.
+ target_sizes (`list[tuple[int, int]]`):
+ list of original (height, width) dimensions for each image before preprocessing.
+ size (`dict[str, int]`):
+ A size dict which was used to resize.
+ """
+ num_classes = segmentation_logits.shape[1]
+ aggregated_logits = []
+ patch_counts = []
+
+ for image_size in target_sizes:
+ height, width = get_size_with_aspect_ratio(image_size, size["shortest_edge"], size["longest_edge"])
+ aggregated_logits.append(torch.zeros((num_classes, height, width), device=segmentation_logits.device))
+ patch_counts.append(torch.zeros((num_classes, height, width), device=segmentation_logits.device))
+
+ # Stitch patches back into full-sized logit maps
+ for patch_idx, (image_idx, patch_start, patch_end) in enumerate(patch_offsets):
+ if target_sizes[image_idx][0] > target_sizes[image_idx][1]:
+ aggregated_logits[image_idx][:, patch_start:patch_end, :] += segmentation_logits[patch_idx]
+ patch_counts[image_idx][:, patch_start:patch_end, :] += 1
+ else:
+ aggregated_logits[image_idx][:, :, patch_start:patch_end] += segmentation_logits[patch_idx]
+ patch_counts[image_idx][:, :, patch_start:patch_end] += 1
+
+ # Normalize and resize logits to original image size
+ reconstructed_logits = []
+ for idx, (logit_sum, count) in enumerate(zip(aggregated_logits, patch_counts)):
+ averaged_logits = logit_sum / count.clamp(min=1)
+ resized_logits = torch.nn.functional.interpolate(
+ averaged_logits[None, ...], size=target_sizes[idx], mode="bilinear", align_corners=False
+ )[0]
+
+ reconstructed_logits.append(resized_logits)
+
+ return reconstructed_logits
+
+ def unpad_image(
+ self, segmentation_logits: "torch.Tensor", target_sizes: list[tuple[int, int]], size: dict[str, int]
+ ) -> "list[torch.Tensor]":
+ """Restores panoptic segmentation logits to their original image resolutions."""
+
+ resized_logits = []
+
+ for idx, original_size in enumerate(target_sizes):
+ target_height, target_width = get_size_with_aspect_ratio(
+ original_size, size["shortest_edge"], size["longest_edge"]
+ )
+ cropped_logits = segmentation_logits[idx][:, :target_height, :target_width]
+ upsampled_logits = torch.nn.functional.interpolate(
+ cropped_logits[None, ...], size=original_size, mode="bilinear", align_corners=False
+ )[0]
+ resized_logits.append(upsampled_logits)
+ return resized_logits
+
+ def post_process_semantic_segmentation(
+ self, outputs, target_sizes: list[tuple[int, int]], size: dict[str, int] | None = None
+ ) -> np.ndarray:
+ """Post-processes model outputs into final semantic segmentation prediction."""
+
+ size = size if size is not None else self.size
+
+ masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width]
+ class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1]
+ patch_offsets = outputs.patch_offsets
+
+ output_size = get_target_size(size)
+ masks_queries_logits = torch.nn.functional.interpolate(masks_queries_logits, size=output_size, mode="bilinear")
+
+ # Remove the null class `[..., :-1]`
+ masks_classes = class_queries_logits.softmax(dim=-1)[..., :-1]
+ masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width]
+
+ segmentation_logits = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs)
+
+ if patch_offsets:
+ output_logits = self.merge_image_patches(segmentation_logits, patch_offsets, target_sizes, size)
+ else:
+ output_logits = []
+
+ for idx in range(len(segmentation_logits)):
+ resized_logits = torch.nn.functional.interpolate(
+ segmentation_logits[idx].unsqueeze(dim=0),
+ size=target_sizes[idx],
+ mode="bilinear",
+ align_corners=False,
+ )
+ output_logits.append(resized_logits[0])
+
+ preds = [logit.argmax(dim=0) for logit in output_logits]
+ return preds
+
+ def post_process_panoptic_segmentation(
+ self,
+ outputs,
+ target_sizes: list[tuple[int, int]],
+ threshold: float = 0.8,
+ mask_threshold: float = 0.5,
+ overlap_mask_area_threshold: float = 0.8,
+ stuff_classes: list[int] | None = None,
+ size: dict[str, int] | None = None,
+ ):
+ """Post-processes model outputs into final panoptic segmentation prediction."""
+
+ size = size if size is not None else self.size
+
+ masks_queries_logits = outputs.masks_queries_logits # [batch_size, num_queries, height, width]
+ class_queries_logits = outputs.class_queries_logits # [batch_size, num_queries, num_classes+1]
+
+ batch_size = class_queries_logits.shape[0]
+ num_labels = class_queries_logits.shape[-1] - 1
+
+ output_size = get_target_size(size)
+ masks_queries_logits = torch.nn.functional.interpolate(masks_queries_logits, size=output_size, mode="bilinear")
+
+ mask_probs_batch = self.unpad_image(masks_queries_logits, target_sizes, size)
+ pred_scores_batch, pred_labels_batch = class_queries_logits.softmax(dim=-1).max(-1)
+
+ results: list = []
+
+ for i in range(batch_size):
+ mask_probs, pred_scores, pred_labels = remove_low_and_no_objects(
+ mask_probs_batch[i], pred_scores_batch[i], pred_labels_batch[i], threshold, num_labels
+ )
+
+ # No mask found
+ if mask_probs.shape[0] <= 0:
+ height, width = target_sizes[i] if target_sizes is not None else mask_probs.shape[1:]
+ segmentation = torch.zeros((height, width)) - 1
+ results.append({"segmentation": segmentation, "segments_info": []})
+ continue
+
+ segmentation, segments = compute_segments(
+ mask_probs=mask_probs,
+ pred_scores=pred_scores,
+ pred_labels=pred_labels,
+ stuff_classes=stuff_classes,
+ mask_threshold=mask_threshold,
+ overlap_mask_area_threshold=overlap_mask_area_threshold,
+ target_size=target_sizes[i] if target_sizes is not None else None,
+ )
+
+ results.append({"segmentation": segmentation, "segments_info": segments})
+ return results
+
+ def post_process_instance_segmentation(
+ self, outputs, target_sizes: list[tuple[int, int]], threshold: float = 0.8, size: dict[str, int] | None = None
+ ):
+ """Post-processes model outputs into Instance Segmentation Predictions."""
+
+ size = size if size is not None else self.size
+
+ masks_queries_logits = outputs.masks_queries_logits
+ class_queries_logits = outputs.class_queries_logits
+
+ output_size = get_target_size(size)
+ masks_queries_logits = torch.nn.functional.interpolate(masks_queries_logits, size=output_size, mode="bilinear")
+
+ mask_probs_batch = self.unpad_image(masks_queries_logits, target_sizes, size)
+
+ device = masks_queries_logits.device
+ batch_size = class_queries_logits.shape[0]
+ num_queries = class_queries_logits.shape[-2]
+
+ results = []
+
+ for i in range(batch_size):
+ mask_pred = mask_probs_batch[i]
+ mask_class = class_queries_logits[i]
+
+ # Remove the null class `[..., :-1]`
+ scores, pred_classes = mask_class.softmax(dim=-1)[..., :-1].max(-1)
+ pred_masks = (mask_pred > 0).float()
+
+ # Calculate average mask prob
+ mask_scores = (mask_pred.sigmoid().flatten(1) * pred_masks.flatten(1)).sum(1) / (
+ pred_masks.flatten(1).sum(1) + 1e-6
+ )
+ pred_scores = scores * mask_scores
+
+ segmentation = torch.zeros(target_sizes[i], device=device) - 1
+
+ instance_maps, segments = [], []
+ current_segment_id = 0
+ for j in range(num_queries):
+ score = pred_scores[j].item()
+
+ if not torch.all(pred_masks[j] == 0) and score >= threshold:
+ segmentation[pred_masks[j] == 1] = current_segment_id
+ segments.append(
+ {
+ "id": current_segment_id,
+ "label_id": pred_classes[j].item(),
+ "score": round(score, 6),
+ }
+ )
+ current_segment_id += 1
+ instance_maps.append(pred_masks[j])
+
+ results.append({"segmentation": segmentation, "segments_info": segments})
+ return results
+
+
+__all__ = ["EomtImageProcessorPil"]
diff --git a/third_party/transformers/src/transformers/models/eomt/modeling_eomt.py b/third_party/transformers/src/transformers/models/eomt/modeling_eomt.py
new file mode 100644
index 0000000000000000000000000000000000000000..a1767132f47ce23653939923db21d3b8f46620d1
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/eomt/modeling_eomt.py
@@ -0,0 +1,1231 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/eomt/modular_eomt.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_eomt.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 Mobile Perception Systems Lab at TU/e and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+
+import collections.abc
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...file_utils import ModelOutput, is_scipy_available, requires_backends
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, is_accelerate_available
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_eomt import EomtConfig
+
+
+if is_scipy_available():
+ from scipy.optimize import linear_sum_assignment
+
+if is_accelerate_available():
+ from accelerate import PartialState
+ from accelerate.utils import reduce
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Class for outputs of [`EomtForUniversalSegmentationOutput`].
+
+ This output can be directly passed to [`~EomtImageProcessor.post_process_semantic_segmentation`] or
+ [`~EomtImageProcessor.post_process_instance_segmentation`] or
+ [`~EomtImageProcessor.post_process_panoptic_segmentation`] to compute final segmentation maps. Please, see
+ [`~EomtImageProcessor] for details regarding usage.
+ """
+)
+class EomtForUniversalSegmentationOutput(ModelOutput):
+ r"""
+ loss (`torch.Tensor`, *optional*):
+ The computed loss, returned when labels are present.
+ class_queries_logits (`torch.FloatTensor`):
+ A tensor of shape `(batch_size, num_queries, num_labels + 1)` representing the proposed classes for each
+ query. Note the `+ 1` is needed because we incorporate the null class.
+ masks_queries_logits (`torch.FloatTensor`):
+ A tensor of shape `(batch_size, num_queries, height, width)` representing the proposed masks for each
+ query.
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Last hidden states (final feature map) of the last layer.
+ 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 + one for the output of each stage) of
+ shape `(batch_size, sequence_length, hidden_size)`. Hidden-states all layers of the model.
+ attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+ Tuple of `tuple(torch.FloatTensor)` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+ sequence_length)`. Self and Cross Attentions weights from transformer decoder.
+ patch_offsets (`list[torch.Tensor]`, *optional*):
+ list of tuples indicating the image index and start and end positions of patches for semantic segmentation.
+ """
+
+ loss: torch.FloatTensor | None = None
+ class_queries_logits: torch.FloatTensor | None = None
+ masks_queries_logits: torch.FloatTensor | None = None
+ last_hidden_state: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ patch_offsets: list[torch.Tensor] | None = None
+
+
+# Adapted from https://github.com/facebookresearch/detectron2/blob/main/projects/PointRend/point_rend/point_features.py
+def sample_point(
+ input_features: torch.Tensor, point_coordinates: torch.Tensor, add_dim=False, **kwargs
+) -> torch.Tensor:
+ """
+ A wrapper around `torch.nn.functional.grid_sample` to support 3D point_coordinates tensors.
+
+ Args:
+ input_features (`torch.Tensor` of shape (batch_size, channels, height, width)):
+ A tensor that contains features map on a height * width grid
+ point_coordinates (`torch.Tensor` of shape (batch_size, num_points, 2) or (batch_size, grid_height, grid_width,:
+ 2)):
+ A tensor that contains [0, 1] * [0, 1] normalized point coordinates
+ add_dim (`bool`):
+ boolean value to keep track of added dimension
+
+ Returns:
+ point_features (`torch.Tensor` of shape (batch_size, channels, num_points) or (batch_size, channels,
+ height_grid, width_grid):
+ A tensor that contains features for points in `point_coordinates`.
+ """
+ if point_coordinates.dim() == 3:
+ add_dim = True
+ point_coordinates = point_coordinates.unsqueeze(2)
+
+ # use nn.function.grid_sample to get features for points in `point_coordinates` via bilinear interpolation
+ point_features = torch.nn.functional.grid_sample(input_features, 2.0 * point_coordinates - 1.0, **kwargs)
+ if add_dim:
+ point_features = point_features.squeeze(3)
+
+ return point_features
+
+
+def pair_wise_dice_loss(inputs: Tensor, labels: Tensor) -> Tensor:
+ """
+ A pair wise version of the dice loss, see `dice_loss` for usage.
+
+ Args:
+ inputs (`torch.Tensor`):
+ A tensor representing a mask
+ labels (`torch.Tensor`):
+ A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+
+ Returns:
+ `torch.Tensor`: The computed loss between each pairs.
+ """
+ inputs = inputs.sigmoid().flatten(1)
+ numerator = 2 * torch.matmul(inputs, labels.T)
+ # using broadcasting to get a [num_queries, NUM_CLASSES] matrix
+ denominator = inputs.sum(-1)[:, None] + labels.sum(-1)[None, :]
+ loss = 1 - (numerator + 1) / (denominator + 1)
+ return loss
+
+
+def pair_wise_sigmoid_cross_entropy_loss(inputs: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+ r"""
+ A pair wise version of the cross entropy loss, see `sigmoid_cross_entropy_loss` for usage.
+
+ Args:
+ inputs (`torch.Tensor`):
+ A tensor representing a mask.
+ labels (`torch.Tensor`):
+ A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+
+ Returns:
+ loss (`torch.Tensor`): The computed loss between each pairs.
+ """
+
+ height_and_width = inputs.shape[1]
+
+ criterion = nn.BCEWithLogitsLoss(reduction="none")
+ cross_entropy_loss_pos = criterion(inputs, torch.ones_like(inputs))
+ cross_entropy_loss_neg = criterion(inputs, torch.zeros_like(inputs))
+
+ loss_pos = torch.matmul(cross_entropy_loss_pos / height_and_width, labels.T)
+ loss_neg = torch.matmul(cross_entropy_loss_neg / height_and_width, (1 - labels).T)
+ loss = loss_pos + loss_neg
+ return loss
+
+
+# Adapted from https://github.com/facebookresearch/Eomt/blob/main/eomt/modeling/matcher.py
+class EomtHungarianMatcher(nn.Module):
+ """This class computes an assignment between the labels and the predictions of the network.
+
+ For efficiency reasons, the labels don't include the no_object. Because of this, in general, there are more
+ predictions than labels. In this case, we do a 1-to-1 matching of the best predictions, while the others are
+ un-matched (and thus treated as non-objects).
+ """
+
+ def __init__(
+ self, cost_class: float = 1.0, cost_mask: float = 1.0, cost_dice: float = 1.0, num_points: int = 12544
+ ):
+ """Creates the matcher
+
+ Params:
+ cost_class (`float`, *optional*, defaults to 1.0):
+ Relative weight of the classification error in the matching cost.
+ cost_mask (`float`, *optional*, defaults to 1.0):
+ This is the relative weight of the focal loss of the binary mask in the matching cost.
+ cost_dice (`float`, *optional*, defaults to 1.0):
+ This is the relative weight of the dice loss of the binary mask in the matching cost.
+ num_points (`int`, *optional*, defaults to 12544):
+ No. of points to sample on which the mask loss will be calculated. The same set of K points are
+ uniformly sampled for all prediction and ground truth masks to construct the cost matrix for bipartite
+ matching.
+ """
+ super().__init__()
+ if cost_class == 0 and cost_mask == 0 and cost_dice == 0:
+ raise ValueError("All costs can't be 0")
+
+ self.num_points = num_points
+ self.cost_class = cost_class
+ self.cost_mask = cost_mask
+ self.cost_dice = cost_dice
+
+ @torch.no_grad()
+ def forward(
+ self,
+ masks_queries_logits: torch.Tensor,
+ class_queries_logits: torch.Tensor,
+ mask_labels: torch.Tensor,
+ class_labels: torch.Tensor,
+ ) -> list[tuple[Tensor]]:
+ """
+ Params:
+ masks_queries_logits (`torch.Tensor`):
+ A tensor of dim `batch_size, num_queries, num_labels` with the classification logits.
+ class_queries_logits (`torch.Tensor`):
+ A tensor of dim `batch_size, num_queries, height, width` with the predicted masks.
+ class_labels (`torch.Tensor`):
+ A tensor of dim `num_target_boxes` (where num_target_boxes is the number of ground-truth objects in the
+ target) containing the class labels.
+ mask_labels (`torch.Tensor`):
+ A tensor of dim `num_target_boxes, height, width` containing the target masks.
+
+ Returns:
+ matched_indices (`list[tuple[Tensor]]`): A list of size batch_size, containing tuples of (index_i, index_j)
+ where:
+ - index_i is the indices of the selected predictions (in order)
+ - index_j is the indices of the corresponding selected labels (in order)
+ For each batch element, it holds:
+ len(index_i) = len(index_j) = min(num_queries, num_target_boxes).
+ """
+ indices: list[tuple[np.array]] = []
+
+ # iterate through batch size
+ batch_size = masks_queries_logits.shape[0]
+ for i in range(batch_size):
+ pred_probs = class_queries_logits[i].softmax(-1)
+ pred_mask = masks_queries_logits[i]
+
+ # Compute the classification cost. Contrary to the loss, we don't use the NLL, but approximate it in 1 - proba[target class]. The 1 is a constant that doesn't change the matching, it can be omitted.
+ cost_class = -pred_probs[:, class_labels[i]]
+ target_mask = mask_labels[i].to(pred_mask)
+ target_mask = target_mask[:, None]
+ pred_mask = pred_mask[:, None]
+
+ # Sample ground truth and predicted masks
+ point_coordinates = torch.rand(1, self.num_points, 2, device=pred_mask.device)
+
+ target_coordinates = point_coordinates.repeat(target_mask.shape[0], 1, 1)
+ target_mask = sample_point(target_mask, target_coordinates, align_corners=False).squeeze(1)
+
+ pred_coordinates = point_coordinates.repeat(pred_mask.shape[0], 1, 1)
+ pred_mask = sample_point(pred_mask, pred_coordinates, align_corners=False).squeeze(1)
+
+ # compute the cross entropy loss between each mask pairs -> shape (num_queries, num_labels)
+ cost_mask = pair_wise_sigmoid_cross_entropy_loss(pred_mask, target_mask)
+ # Compute the dice loss between each mask pairs -> shape (num_queries, num_labels)
+ cost_dice = pair_wise_dice_loss(pred_mask, target_mask)
+ # final cost matrix
+ cost_matrix = self.cost_mask * cost_mask + self.cost_class * cost_class + self.cost_dice * cost_dice
+ # eliminate infinite values in cost_matrix to avoid the error ``ValueError: cost matrix is infeasible``
+ cost_matrix = torch.minimum(cost_matrix, torch.tensor(1e10))
+ cost_matrix = torch.maximum(cost_matrix, torch.tensor(-1e10))
+ cost_matrix = torch.nan_to_num(cost_matrix, 0)
+ # do the assignment using the hungarian algorithm in scipy
+ assigned_indices: tuple[np.array] = linear_sum_assignment(cost_matrix.cpu())
+ indices.append(assigned_indices)
+
+ # It could be stacked in one tensor
+ matched_indices = [
+ (torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices
+ ]
+ return matched_indices
+
+
+def dice_loss(inputs: Tensor, labels: Tensor, num_masks: int) -> Tensor:
+ r"""
+ Compute the DICE loss, similar to generalized IOU for masks as follows:
+
+ $$ \mathcal{L}_{\text{dice}(x, y) = 1 - \frac{2 * x \cap y }{x \cup y + 1}} $$
+
+ In practice, since `labels` is a binary mask, (only 0s and 1s), dice can be computed as follow
+
+ $$ \mathcal{L}_{\text{dice}(x, y) = 1 - \frac{2 * x * y }{x + y + 1}} $$
+
+ Args:
+ inputs (`torch.Tensor`):
+ A tensor representing a mask.
+ labels (`torch.Tensor`):
+ A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+ num_masks (`int`):
+ The number of masks present in the current batch, used for normalization.
+
+ Returns:
+ `torch.Tensor`: The computed loss.
+ """
+ probs = inputs.sigmoid().flatten(1)
+ numerator = 2 * (probs * labels).sum(-1)
+ denominator = probs.sum(-1) + labels.sum(-1)
+ loss = 1 - (numerator + 1) / (denominator + 1)
+ loss = loss.sum() / num_masks
+ return loss
+
+
+def sigmoid_cross_entropy_loss(inputs: torch.Tensor, labels: torch.Tensor, num_masks: int) -> torch.Tensor:
+ r"""
+ Args:
+ inputs (`torch.Tensor`):
+ A float tensor of arbitrary shape.
+ labels (`torch.Tensor`):
+ A tensor with the same shape as inputs. Stores the binary classification labels for each element in inputs
+ (0 for the negative class and 1 for the positive class).
+
+ Returns:
+ loss (`torch.Tensor`): The computed loss.
+ """
+ criterion = nn.BCEWithLogitsLoss(reduction="none")
+ cross_entropy_loss = criterion(inputs, labels)
+
+ loss = cross_entropy_loss.mean(1).sum() / num_masks
+ return loss
+
+
+# Adapted from https://github.com/facebookresearch/Eomt/blob/main/eomt/modeling/criterion.py
+class EomtLoss(nn.Module):
+ def __init__(self, config: EomtConfig, weight_dict: dict[str, float]):
+ """
+ The Eomt Loss. The loss is computed very similar to DETR. The process happens in two steps: 1) we
+ compute hungarian assignment between ground truth masks and the outputs of the model 2) we supervise each pair
+ of matched ground-truth / prediction (supervise class and mask)
+
+ Args:
+ config (`EomtConfig`):
+ The configuration for Eomt model also containing loss calculation specific parameters.
+ weight_dict (`dict[str, float]`):
+ A dictionary of weights to be applied to the different losses.
+ """
+ super().__init__()
+ requires_backends(self, ["scipy"])
+ self.num_labels = config.num_labels
+ self.weight_dict = weight_dict
+
+ # Weight to apply to the null class
+ self.eos_coef = config.no_object_weight
+ empty_weight = torch.ones(self.num_labels + 1)
+ empty_weight[-1] = self.eos_coef
+ self.register_buffer("empty_weight", empty_weight)
+
+ # pointwise mask loss parameters
+ self.num_points = config.train_num_points
+ self.oversample_ratio = config.oversample_ratio
+ self.importance_sample_ratio = config.importance_sample_ratio
+
+ self.matcher = EomtHungarianMatcher(
+ cost_class=config.class_weight,
+ cost_dice=config.dice_weight,
+ cost_mask=config.mask_weight,
+ num_points=self.num_points,
+ )
+
+ def _max_by_axis(self, sizes: list[list[int]]) -> list[int]:
+ maxes = sizes[0]
+ for sublist in sizes[1:]:
+ for index, item in enumerate(sublist):
+ maxes[index] = max(maxes[index], item)
+ return maxes
+
+ # Adapted from nested_tensor_from_tensor_list() in original implementation
+ def _pad_images_to_max_in_batch(self, tensors: list[Tensor]) -> tuple[Tensor, Tensor]:
+ # get the maximum size in the batch
+ max_size = self._max_by_axis([list(tensor.shape) for tensor in tensors])
+ # compute final size
+ batch_shape = [len(tensors)] + max_size
+ batch_size, _, height, width = batch_shape
+ dtype = tensors[0].dtype
+ device = tensors[0].device
+ padded_tensors = torch.zeros(batch_shape, dtype=dtype, device=device)
+ padding_masks = torch.ones((batch_size, height, width), dtype=torch.bool, device=device)
+ # pad the tensors to the size of the biggest one
+ for tensor, padded_tensor, padding_mask in zip(tensors, padded_tensors, padding_masks):
+ padded_tensor[: tensor.shape[0], : tensor.shape[1], : tensor.shape[2]].copy_(tensor)
+ padding_mask[: tensor.shape[1], : tensor.shape[2]] = False
+
+ return padded_tensors, padding_masks
+
+ def loss_labels(
+ self, class_queries_logits: Tensor, class_labels: list[Tensor], indices: tuple[np.array]
+ ) -> dict[str, Tensor]:
+ """Compute the losses related to the labels using cross entropy.
+
+ Args:
+ class_queries_logits (`torch.Tensor`):
+ A tensor of shape `batch_size, num_queries, num_labels`
+ class_labels (`list[torch.Tensor]`):
+ List of class labels of shape `(labels)`.
+ indices (`tuple[np.array])`:
+ The indices computed by the Hungarian matcher.
+
+ Returns:
+ `dict[str, Tensor]`: A dict of `torch.Tensor` containing the following key:
+ - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels.
+ """
+ pred_logits = class_queries_logits
+ batch_size, num_queries, _ = pred_logits.shape
+ criterion = nn.CrossEntropyLoss(weight=self.empty_weight)
+ idx = self._get_predictions_permutation_indices(indices) # shape of (batch_size, num_queries)
+ target_classes_o = torch.cat(
+ [target[j] for target, (_, j) in zip(class_labels, indices)]
+ ) # shape of (batch_size, num_queries)
+ target_classes = torch.full(
+ (batch_size, num_queries), fill_value=self.num_labels, dtype=torch.int64, device=pred_logits.device
+ )
+ target_classes[idx] = target_classes_o
+ # Permute target_classes (batch_size, num_queries, num_labels) -> (batch_size, num_labels, num_queries)
+ pred_logits_transposed = pred_logits.transpose(1, 2)
+ loss_ce = criterion(pred_logits_transposed, target_classes)
+ losses = {"loss_cross_entropy": loss_ce}
+ return losses
+
+ def loss_masks(
+ self,
+ masks_queries_logits: torch.Tensor,
+ mask_labels: list[torch.Tensor],
+ indices: tuple[np.array],
+ num_masks: int,
+ ) -> dict[str, torch.Tensor]:
+ """Compute the losses related to the masks using sigmoid_cross_entropy_loss and dice loss.
+
+ Args:
+ masks_queries_logits (`torch.Tensor`):
+ A tensor of shape `(batch_size, num_queries, height, width)`.
+ mask_labels (`torch.Tensor`):
+ List of mask labels of shape `(labels, height, width)`.
+ indices (`tuple[np.array])`:
+ The indices computed by the Hungarian matcher.
+ num_masks (`int)`:
+ The number of masks, used for normalization.
+
+ Returns:
+ losses (`dict[str, Tensor]`): A dict of `torch.Tensor` containing two keys:
+ - **loss_mask** -- The loss computed using sigmoid cross entropy loss on the predicted and ground truth.
+ masks.
+ - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth,
+ masks.
+ """
+ src_idx = self._get_predictions_permutation_indices(indices)
+ tgt_idx = self._get_targets_permutation_indices(indices)
+ # shape (batch_size * num_queries, height, width)
+ pred_masks = masks_queries_logits[src_idx]
+ # shape (batch_size, num_queries, height, width)
+ # pad all and stack the targets to the num_labels dimension
+ target_masks, _ = self._pad_images_to_max_in_batch(mask_labels)
+ target_masks = target_masks[tgt_idx]
+
+ # No need to upsample predictions as we are using normalized coordinates
+ pred_masks = pred_masks[:, None]
+ target_masks = target_masks[:, None]
+
+ # Sample point coordinates
+ with torch.no_grad():
+ point_coordinates = self.sample_points_using_uncertainty(
+ pred_masks,
+ lambda logits: self.calculate_uncertainty(logits),
+ self.num_points,
+ self.oversample_ratio,
+ self.importance_sample_ratio,
+ )
+
+ point_labels = sample_point(target_masks, point_coordinates, align_corners=False).squeeze(1)
+
+ point_logits = sample_point(pred_masks, point_coordinates, align_corners=False).squeeze(1)
+
+ losses = {
+ "loss_mask": sigmoid_cross_entropy_loss(point_logits, point_labels, num_masks),
+ "loss_dice": dice_loss(point_logits, point_labels, num_masks),
+ }
+
+ del pred_masks
+ del target_masks
+ return losses
+
+ def _get_predictions_permutation_indices(self, indices):
+ # Permute predictions following indices
+ batch_indices = torch.cat([torch.full_like(src, i) for i, (src, _) in enumerate(indices)])
+ predictions_indices = torch.cat([src for (src, _) in indices])
+ return batch_indices, predictions_indices
+
+ def _get_targets_permutation_indices(self, indices):
+ # Permute labels following indices
+ batch_indices = torch.cat([torch.full_like(tgt, i) for i, (_, tgt) in enumerate(indices)])
+ target_indices = torch.cat([tgt for (_, tgt) in indices])
+ return batch_indices, target_indices
+
+ def calculate_uncertainty(self, logits: torch.Tensor) -> torch.Tensor:
+ """
+ In Eomt paper, uncertainty is estimated as L1 distance between 0.0 and the logit prediction in 'logits'
+ for the foreground class in `classes`.
+
+ Args:
+ logits (`torch.Tensor`):
+ A tensor of shape (R, 1, ...) for class-specific or class-agnostic, where R is the total number of predicted masks in all images and C is:
+ the number of foreground classes. The values are logits.
+
+ Returns:
+ scores (`torch.Tensor`): A tensor of shape (R, 1, ...) that contains uncertainty scores with the most
+ uncertain locations having the highest uncertainty score.
+ """
+ uncertainty_scores = -(torch.abs(logits))
+ return uncertainty_scores
+
+ def sample_points_using_uncertainty(
+ self,
+ logits: torch.Tensor,
+ uncertainty_function,
+ num_points: int,
+ oversample_ratio: int,
+ importance_sample_ratio: float,
+ ) -> torch.Tensor:
+ """
+ This function is meant for sampling points in [0, 1] * [0, 1] coordinate space based on their uncertainty. The
+ uncertainty is calculated for each point using the passed `uncertainty function` that takes points logit
+ prediction as input.
+
+ Args:
+ logits (`float`):
+ Logit predictions for P points.
+ uncertainty_function:
+ A function that takes logit predictions for P points and returns their uncertainties.
+ num_points (`int`):
+ The number of points P to sample.
+ oversample_ratio (`int`):
+ Oversampling parameter.
+ importance_sample_ratio (`float`):
+ Ratio of points that are sampled via importance sampling.
+
+ Returns:
+ point_coordinates (`torch.Tensor`):
+ Coordinates for P sampled points.
+ """
+
+ num_boxes = logits.shape[0]
+ num_points_sampled = int(num_points * oversample_ratio)
+
+ # Get random point coordinates
+ point_coordinates = torch.rand(num_boxes, num_points_sampled, 2, device=logits.device)
+ # Get sampled prediction value for the point coordinates
+ point_logits = sample_point(logits, point_coordinates, align_corners=False)
+ # Calculate the uncertainties based on the sampled prediction values of the points
+ point_uncertainties = uncertainty_function(point_logits)
+
+ num_uncertain_points = int(importance_sample_ratio * num_points)
+ num_random_points = num_points - num_uncertain_points
+
+ idx = torch.topk(point_uncertainties[:, 0, :], k=num_uncertain_points, dim=1)[1]
+ shift = num_points_sampled * torch.arange(num_boxes, dtype=torch.long, device=logits.device)
+ idx += shift[:, None]
+ point_coordinates = point_coordinates.view(-1, 2)[idx.view(-1), :].view(num_boxes, num_uncertain_points, 2)
+
+ if num_random_points > 0:
+ point_coordinates = torch.cat(
+ [point_coordinates, torch.rand(num_boxes, num_random_points, 2, device=logits.device)],
+ dim=1,
+ )
+ return point_coordinates
+
+ def forward(
+ self,
+ masks_queries_logits: torch.Tensor,
+ class_queries_logits: torch.Tensor,
+ mask_labels: list[torch.Tensor],
+ class_labels: list[torch.Tensor],
+ auxiliary_predictions: dict[str, torch.Tensor] | None = None,
+ ) -> dict[str, torch.Tensor]:
+ """
+ This performs the loss computation.
+
+ Args:
+ masks_queries_logits (`torch.Tensor`):
+ A tensor of shape `(batch_size, num_queries, height, width)`.
+ class_queries_logits (`torch.Tensor`):
+ A tensor of shape `(batch_size, num_queries, num_labels)`.
+ mask_labels (`torch.Tensor`):
+ List of mask labels of shape `(labels, height, width)`.
+ class_labels (`list[torch.Tensor]`):
+ List of class labels of shape `(labels)`.
+ auxiliary_predictions (`dict[str, torch.Tensor]`, *optional*):
+ if `use_auxiliary_loss` was set to `true` in [`EomtConfig`], then it contains the logits from
+ the inner layers of the EomtMaskedAttentionDecoder.
+
+ Returns:
+ losses (`dict[str, Tensor]`): A dict of `torch.Tensor` containing three keys:
+ - **loss_cross_entropy** -- The loss computed using cross entropy on the predicted and ground truth labels.
+ - **loss_mask** -- The loss computed using sigmoid cross_entropy loss on the predicted and ground truth
+ masks.
+ - **loss_dice** -- The loss computed using dice loss on the predicted on the predicted and ground truth
+ masks.
+ if `use_auxiliary_loss` was set to `true` in [`EomtConfig`], the dictionary contains additional
+ losses for each auxiliary predictions.
+ """
+
+ # retrieve the matching between the outputs of the last layer and the labels
+ indices = self.matcher(masks_queries_logits, class_queries_logits, mask_labels, class_labels)
+ # compute the average number of target masks for normalization purposes
+ num_masks = self.get_num_masks(class_labels, device=class_labels[0].device)
+ # get all the losses
+ losses: dict[str, Tensor] = {
+ **self.loss_masks(masks_queries_logits, mask_labels, indices, num_masks),
+ **self.loss_labels(class_queries_logits, class_labels, indices),
+ }
+ # in case of auxiliary losses, we repeat this process with the output of each intermediate layer.
+ if auxiliary_predictions is not None:
+ for idx, aux_outputs in enumerate(auxiliary_predictions):
+ masks_queries_logits = aux_outputs["masks_queries_logits"]
+ class_queries_logits = aux_outputs["class_queries_logits"]
+ loss_dict = self.forward(masks_queries_logits, class_queries_logits, mask_labels, class_labels)
+ loss_dict = {f"{key}_{idx}": value for key, value in loss_dict.items()}
+ losses.update(loss_dict)
+
+ return losses
+
+ def get_num_masks(self, class_labels: torch.Tensor, device: torch.device) -> torch.Tensor:
+ """
+ Computes the average number of target masks across the batch, for normalization purposes.
+ """
+ num_masks = sum(len(classes) for classes in class_labels)
+ num_masks = torch.as_tensor(num_masks, dtype=torch.float, device=device)
+ world_size = 1
+ if is_accelerate_available():
+ if PartialState._shared_state != {}:
+ num_masks = reduce(num_masks)
+ world_size = PartialState().num_processes
+
+ num_masks = torch.clamp(num_masks / world_size, min=1)
+ return num_masks
+
+
+class EomtPatchEmbeddings(nn.Module):
+ """
+ This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
+ `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
+ Transformer.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ image_size, patch_size = config.image_size, config.patch_size
+ num_channels, hidden_size = config.num_channels, config.hidden_size
+
+ image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
+ patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
+ num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
+ self.image_size = image_size
+ self.patch_size = patch_size
+ self.num_channels = num_channels
+ self.num_patches = num_patches
+
+ self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)
+
+ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+ num_channels = pixel_values.shape[1]
+ if num_channels != self.num_channels:
+ raise ValueError(
+ "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+ f" Expected {self.num_channels} but got {num_channels}."
+ )
+ embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2)
+ return embeddings
+
+
+class EomtEmbeddings(nn.Module):
+ """
+ Construct the CLS token, mask token, position and patch embeddings.
+ """
+
+ def __init__(self, config: EomtConfig) -> None:
+ super().__init__()
+
+ self.config = config
+ self.patch_size = config.patch_size
+
+ self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_size))
+ self.register_tokens = nn.Parameter(torch.zeros(1, config.num_register_tokens, config.hidden_size))
+
+ self.patch_embeddings = EomtPatchEmbeddings(config)
+ num_patches = self.patch_embeddings.num_patches
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ self.num_prefix_tokens = 1 + config.num_register_tokens # 1 for [CLS]
+ self.position_embeddings = nn.Embedding(num_patches, config.hidden_size)
+ self.register_buffer("position_ids", torch.arange(num_patches).expand((1, -1)), persistent=False)
+
+ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+ batch_size, _, _, _ = pixel_values.shape
+ target_dtype = self.patch_embeddings.projection.weight.dtype
+ embeddings = self.patch_embeddings(pixel_values.to(dtype=target_dtype))
+
+ cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+ register_tokens = self.register_tokens.expand(batch_size, -1, -1)
+
+ embeddings = embeddings + self.position_embeddings(self.position_ids)
+ embeddings = torch.cat([cls_tokens, register_tokens, embeddings], dim=1)
+
+ embeddings = self.dropout(embeddings)
+
+ return embeddings
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class EomtAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = self.embed_dim // self.num_heads
+ if self.head_dim * self.num_heads != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+ f" {self.num_heads})."
+ )
+ self.scale = self.head_dim**-0.5
+ self.dropout = config.attention_dropout
+ self.is_causal = False
+
+ self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Input shape: Batch x Time x Channel"""
+
+ batch_size, seq_length, embed_dim = hidden_states.shape
+
+ queries = self.q_proj(hidden_states)
+ keys = self.k_proj(hidden_states)
+ values = self.v_proj(hidden_states)
+
+ queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+ keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+ values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ queries,
+ keys,
+ values,
+ attention_mask,
+ is_causal=self.is_causal,
+ scaling=self.scale,
+ dropout=0.0 if not self.training else self.dropout,
+ )
+
+ attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous()
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+class EomtLayerScale(nn.Module):
+ def __init__(self, config) -> None:
+ super().__init__()
+ self.lambda1 = nn.Parameter(config.layerscale_value * torch.ones(config.hidden_size))
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ return hidden_state * self.lambda1
+
+
+def drop_path(input: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
+ """
+ Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
+
+ """
+ if drop_prob == 0.0 or not training:
+ return input
+ keep_prob = 1 - drop_prob
+ shape = (input.shape[0],) + (1,) * (input.ndim - 1) # work with diff dim tensors, not just 2D ConvNets
+ random_tensor = keep_prob + torch.rand(shape, dtype=input.dtype, device=input.device)
+ random_tensor.floor_() # binarize
+ output = input.div(keep_prob) * random_tensor
+ return output
+
+
+class EomtDropPath(nn.Module):
+ """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
+
+ def __init__(self, drop_prob: float | None = None) -> None:
+ super().__init__()
+ self.drop_prob = drop_prob
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ return drop_path(hidden_states, self.drop_prob, self.training)
+
+ def extra_repr(self) -> str:
+ return f"p={self.drop_prob}"
+
+
+class EomtMLP(nn.Module):
+ def __init__(self, config) -> None:
+ super().__init__()
+ in_features = out_features = config.hidden_size
+ hidden_features = int(config.hidden_size * config.mlp_ratio)
+ self.fc1 = nn.Linear(in_features, hidden_features, bias=True)
+ if isinstance(config.hidden_act, str):
+ self.activation = ACT2FN[config.hidden_act]
+ else:
+ self.activation = config.hidden_act
+ self.fc2 = nn.Linear(hidden_features, out_features, bias=True)
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ hidden_state = self.fc1(hidden_state)
+ hidden_state = self.activation(hidden_state)
+ hidden_state = self.fc2(hidden_state)
+ return hidden_state
+
+
+class EomtSwiGLUFFN(nn.Module):
+ def __init__(self, config) -> None:
+ super().__init__()
+ in_features = out_features = config.hidden_size
+ hidden_features = int(config.hidden_size * config.mlp_ratio)
+ hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8
+
+ self.weights_in = nn.Linear(in_features, 2 * hidden_features, bias=True)
+ self.weights_out = nn.Linear(hidden_features, out_features, bias=True)
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ hidden_state = self.weights_in(hidden_state)
+ x1, x2 = hidden_state.chunk(2, dim=-1)
+ hidden = nn.functional.silu(x1) * x2
+ return self.weights_out(hidden)
+
+
+class EomtLayer(GradientCheckpointingLayer):
+ """This corresponds to the Block class in the original implementation."""
+
+ def __init__(self, config: EomtConfig) -> None:
+ super().__init__()
+
+ self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.attention = EomtAttention(config)
+ self.layer_scale1 = EomtLayerScale(config)
+ self.drop_path = EomtDropPath(config.drop_path_rate) if config.drop_path_rate > 0.0 else nn.Identity()
+
+ self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+ if config.use_swiglu_ffn:
+ self.mlp = EomtSwiGLUFFN(config)
+ else:
+ self.mlp = EomtMLP(config)
+ self.layer_scale2 = EomtLayerScale(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ hidden_states_norm = self.norm1(hidden_states)
+ self_attention_output, _ = self.attention(hidden_states_norm, attention_mask)
+ self_attention_output = self.layer_scale1(self_attention_output)
+
+ # first residual connection
+ hidden_states = self.drop_path(self_attention_output) + hidden_states
+
+ # in Eomt, layernorm is also applied after self-attention
+ layer_output = self.norm2(hidden_states)
+ layer_output = self.mlp(layer_output)
+ layer_output = self.layer_scale2(layer_output)
+
+ # second residual connection
+ layer_output = self.drop_path(layer_output) + hidden_states
+
+ return layer_output
+
+
+class EomtLayerNorm2d(nn.LayerNorm):
+ def __init__(self, num_channels, eps=1e-6, affine=True):
+ super().__init__(num_channels, eps=eps, elementwise_affine=affine)
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ hidden_state = hidden_state.permute(0, 2, 3, 1)
+ hidden_state = F.layer_norm(hidden_state, self.normalized_shape, self.weight, self.bias, self.eps)
+ hidden_state = hidden_state.permute(0, 3, 1, 2)
+ return hidden_state
+
+
+class EomtScaleLayer(nn.Module):
+ def __init__(self, config: EomtConfig):
+ super().__init__()
+ hidden_size = config.hidden_size
+ self.conv1 = nn.ConvTranspose2d(hidden_size, hidden_size, kernel_size=2, stride=2)
+ self.activation = ACT2FN[config.hidden_act]
+ self.conv2 = nn.Conv2d(
+ hidden_size,
+ hidden_size,
+ kernel_size=3,
+ padding=1,
+ groups=hidden_size,
+ bias=False,
+ )
+
+ self.layernorm2d = EomtLayerNorm2d(hidden_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.conv1(hidden_states)
+ hidden_states = self.activation(hidden_states)
+ hidden_states = self.conv2(hidden_states)
+ hidden_states = self.layernorm2d(hidden_states)
+ return hidden_states
+
+
+class EomtScaleBlock(nn.Module):
+ def __init__(self, config: EomtConfig):
+ super().__init__()
+ self.num_blocks = config.num_upscale_blocks
+ self.block = nn.ModuleList([EomtScaleLayer(config) for _ in range(self.num_blocks)])
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ for block in self.block:
+ hidden_states = block(hidden_states)
+ return hidden_states
+
+
+class EomtMaskHead(nn.Module):
+ def __init__(self, config: EomtConfig):
+ super().__init__()
+
+ hidden_size = config.hidden_size
+ self.fc1 = nn.Linear(hidden_size, hidden_size)
+ self.fc2 = nn.Linear(hidden_size, hidden_size)
+ self.fc3 = nn.Linear(hidden_size, hidden_size)
+ self.activation = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.activation(self.fc1(hidden_states))
+ hidden_states = self.activation(self.fc2(hidden_states))
+ hidden_states = self.fc3(hidden_states)
+ return hidden_states
+
+
+@auto_docstring
+class EomtPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config: EomtConfig
+ base_model_prefix = "eomt"
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+ supports_gradient_checkpointing = False
+ _no_split_modules = ["EomtLayer"]
+ _supports_sdpa = True
+ _can_record_outputs = {
+ "hidden_states": EomtLayer,
+ "attentions": EomtAttention,
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module: nn.Module) -> None:
+ std = self.config.initializer_range
+ if isinstance(module, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)):
+ init.kaiming_uniform_(module.weight, a=math.sqrt(5))
+ if module.bias is not None:
+ fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(module.weight)
+ bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+ init.uniform_(module.bias, -bound, bound)
+ elif isinstance(module, nn.LayerNorm):
+ init.ones_(module.weight)
+ init.zeros_(module.bias)
+ elif isinstance(module, nn.Embedding):
+ init.normal_(module.weight, mean=0.0, std=1)
+ # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag
+ if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False):
+ init.zeros_(module.weight[module.padding_idx])
+ elif isinstance(module, EomtLayerScale):
+ if hasattr(module, "lambda1"):
+ init.constant_(module.lambda1, self.config.layerscale_value)
+ elif isinstance(module, EomtEmbeddings):
+ init.trunc_normal_(module.cls_token, mean=0.0, std=std)
+ init.zeros_(module.register_tokens)
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+ elif isinstance(module, EomtLoss):
+ empty_weight = torch.ones(module.num_labels + 1)
+ empty_weight[-1] = module.eos_coef
+ init.copy_(module.empty_weight, empty_weight)
+ elif isinstance(module, EomtForUniversalSegmentation):
+ init.ones_(module.attn_mask_probs)
+
+
+@auto_docstring(
+ custom_intro="""
+ The EoMT Model with head on top for instance/semantic/panoptic segmentation.
+ """
+)
+class EomtForUniversalSegmentation(EomtPreTrainedModel):
+ main_input_name = "pixel_values"
+
+ def __init__(self, config: EomtConfig):
+ super().__init__(config)
+ self.config = config
+ self.num_hidden_layers = config.num_hidden_layers
+ self.embeddings = EomtEmbeddings(config)
+ self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+ self.query = nn.Embedding(config.num_queries, config.hidden_size)
+ self.layers = nn.ModuleList([EomtLayer(config) for _ in range(config.num_hidden_layers)])
+
+ self.upscale_block = EomtScaleBlock(config)
+ self.mask_head = EomtMaskHead(config)
+
+ self.class_predictor = nn.Linear(config.hidden_size, config.num_labels + 1)
+
+ self.grid_size = (config.image_size // config.patch_size, config.image_size // config.patch_size)
+ self.weight_dict: dict[str, float] = {
+ "loss_cross_entropy": config.class_weight,
+ "loss_mask": config.mask_weight,
+ "loss_dice": config.dice_weight,
+ }
+
+ self.criterion = EomtLoss(config=config, weight_dict=self.weight_dict)
+
+ self.register_buffer("attn_mask_probs", torch.ones(config.num_blocks))
+
+ self.post_init()
+
+ def get_loss_dict(
+ self,
+ masks_queries_logits: Tensor,
+ class_queries_logits: Tensor,
+ mask_labels: Tensor,
+ class_labels: Tensor,
+ auxiliary_predictions: dict[str, Tensor],
+ ) -> dict[str, Tensor]:
+ loss_dict: dict[str, Tensor] = self.criterion(
+ masks_queries_logits=masks_queries_logits,
+ class_queries_logits=class_queries_logits,
+ mask_labels=mask_labels,
+ class_labels=class_labels,
+ auxiliary_predictions=auxiliary_predictions,
+ )
+
+ # weight each loss by `self.weight_dict[]` including auxiliary losses
+ for key, weight in self.weight_dict.items():
+ for loss_key, loss in loss_dict.items():
+ if key in loss_key:
+ loss *= weight
+
+ return loss_dict
+
+ def get_loss(self, loss_dict: dict[str, Tensor]) -> Tensor:
+ return sum(loss_dict.values())
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: Tensor,
+ mask_labels: list[Tensor] | None = None,
+ class_labels: list[Tensor] | None = None,
+ patch_offsets: list[Tensor] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> EomtForUniversalSegmentationOutput:
+ r"""
+ mask_labels (`list[torch.Tensor]`, *optional*):
+ list of mask labels of shape `(num_labels, height, width)` to be fed to a model
+ class_labels (`list[torch.LongTensor]`, *optional*):
+ list of target class labels of shape `(num_labels, height, width)` to be fed to a model. They identify the
+ labels of `mask_labels`, e.g. the label of `mask_labels[i][j]` if `class_labels[i][j]`.
+ patch_offsets (`list[torch.Tensor]`, *optional*):
+ list of tuples indicating the image index and start and end positions of patches for semantic segmentation.
+ """
+
+ masks_queries_logits_per_layer, class_queries_logits_per_layer = (), ()
+ attention_mask = None
+
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ hidden_states = self.embeddings(pixel_values)
+
+ for idx, layer_module in enumerate(self.layers):
+ if idx == self.num_hidden_layers - self.config.num_blocks:
+ query = self.query.weight[None, :, :].expand(hidden_states.shape[0], -1, -1).to(hidden_states.device)
+ hidden_states = torch.cat((query, hidden_states), dim=1)
+
+ if idx >= self.num_hidden_layers - self.config.num_blocks and (
+ self.training or self.attn_mask_probs[idx - self.num_hidden_layers + self.config.num_blocks] > 0
+ ):
+ norm_hidden_states = self.layernorm(hidden_states)
+ masks_queries_logits, class_queries_logits = self.predict(norm_hidden_states)
+
+ masks_queries_logits_per_layer += (masks_queries_logits,)
+ class_queries_logits_per_layer += (class_queries_logits,)
+
+ attention_mask = torch.ones(
+ hidden_states.shape[0],
+ hidden_states.shape[1],
+ hidden_states.shape[1],
+ device=hidden_states.device,
+ dtype=torch.bool,
+ )
+
+ interpolated_logits = F.interpolate(masks_queries_logits, size=self.grid_size, mode="bilinear")
+ interpolated_logits = interpolated_logits.view(
+ interpolated_logits.size(0), interpolated_logits.size(1), -1
+ )
+
+ num_query_tokens = self.config.num_queries
+ encoder_start_tokens = num_query_tokens + self.embeddings.num_prefix_tokens
+
+ # Set attention mask for queries to focus on encoder tokens based on interpolated logits
+ attention_mask[:, :num_query_tokens, encoder_start_tokens:] = interpolated_logits > 0
+
+ # Disable attention mask for random query tokens.
+ attention_mask = self._disable_attention_mask(
+ attention_mask,
+ prob=self.attn_mask_probs[idx - self.num_hidden_layers + self.config.num_blocks],
+ num_query_tokens=num_query_tokens,
+ encoder_start_tokens=encoder_start_tokens,
+ device=attention_mask.device,
+ )
+
+ # Expand attention mask to 4d mask.
+ attention_mask = attention_mask[:, None, ...].expand(-1, self.config.num_attention_heads, -1, -1)
+ attention_mask = attention_mask.float().masked_fill(~attention_mask, -1e9)
+
+ hidden_states = layer_module(hidden_states, attention_mask)
+
+ sequence_output = self.layernorm(hidden_states)
+
+ masks_queries_logits, class_queries_logits = self.predict(sequence_output)
+ masks_queries_logits_per_layer += (masks_queries_logits,)
+ class_queries_logits_per_layer += (class_queries_logits,)
+
+ loss = None
+ if mask_labels is not None and class_labels is not None:
+ loss = 0.0
+ for masks_queries_logits, class_queries_logits in zip(
+ masks_queries_logits_per_layer, class_queries_logits_per_layer
+ ):
+ loss_dict = self.get_loss_dict(
+ masks_queries_logits=masks_queries_logits,
+ class_queries_logits=class_queries_logits,
+ mask_labels=mask_labels,
+ class_labels=class_labels,
+ auxiliary_predictions=None,
+ )
+ loss += self.get_loss(loss_dict)
+
+ return EomtForUniversalSegmentationOutput(
+ loss=loss,
+ masks_queries_logits=masks_queries_logits,
+ class_queries_logits=class_queries_logits,
+ last_hidden_state=sequence_output,
+ patch_offsets=patch_offsets,
+ )
+
+ def get_input_embeddings(self):
+ return self.embeddings.patch_embeddings
+
+ def predict(self, logits: torch.Tensor):
+ query_tokens = logits[:, : self.config.num_queries, :]
+ class_logits = self.class_predictor(query_tokens)
+
+ prefix_tokens = logits[:, self.config.num_queries + self.embeddings.num_prefix_tokens :, :]
+ prefix_tokens = prefix_tokens.transpose(1, 2)
+
+ prefix_tokens = prefix_tokens.reshape(prefix_tokens.shape[0], -1, *self.grid_size)
+
+ query_tokens = self.mask_head(query_tokens)
+ prefix_tokens = self.upscale_block(prefix_tokens)
+
+ mask_logits = torch.einsum("bqc, bchw -> bqhw", query_tokens, prefix_tokens)
+
+ return mask_logits, class_logits
+
+ @staticmethod
+ def _disable_attention_mask(attn_mask, prob, num_query_tokens, encoder_start_tokens, device):
+ if prob < 1:
+ # Generate random queries to disable based on the probs
+ random_queries = torch.rand(attn_mask.shape[0], num_query_tokens, device=device) > prob
+
+ # Disable attention to the query tokens, considering the prefix tokens
+ attn_mask[:, :num_query_tokens, encoder_start_tokens:][random_queries] = 1
+
+ return attn_mask
+
+
+__all__ = ["EomtPreTrainedModel", "EomtForUniversalSegmentation"]
diff --git a/third_party/transformers/src/transformers/models/eomt/modular_eomt.py b/third_party/transformers/src/transformers/models/eomt/modular_eomt.py
new file mode 100644
index 0000000000000000000000000000000000000000..e4dafa024861653be6d0d53b237b0ba92df6a211
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/eomt/modular_eomt.py
@@ -0,0 +1,544 @@
+# Copyright 2025 Mobile Perception Systems Lab at TU/e and The HuggingFace Inc. team. All rights reserved.
+#
+# 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 EoMT model."""
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from huggingface_hub.dataclasses import strict
+from torch import Tensor, nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...file_utils import (
+ ModelOutput,
+)
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import (
+ TransformersKwargs,
+ auto_docstring,
+ logging,
+)
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from ..dinov2.modeling_dinov2 import (
+ Dinov2Embeddings,
+ Dinov2Layer,
+ Dinov2LayerScale,
+ Dinov2PatchEmbeddings,
+)
+from ..mask2former.modeling_mask2former import Mask2FormerForUniversalSegmentation, Mask2FormerLoss
+from ..siglip.modeling_siglip import SiglipAttention
+from ..vit.configuration_vit import ViTConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="tue-mps/coco_panoptic_eomt_large_640")
+@strict
+class EomtConfig(ViTConfig):
+ r"""
+ layerscale_value (`float`, *optional*, defaults to 1.0):
+ Initial value for the LayerScale parameter.
+ num_upscale_blocks (`int`, *optional*, defaults to 2):
+ Number of upsampling blocks used in the decoder or segmentation head.
+ use_swiglu_ffn (`bool`, *optional*, defaults to `False`):
+ Whether to use the SwiGLU feedforward neural network.
+ num_blocks (`int`, *optional*, defaults to 4):
+ Number of feature blocks or stages in the architecture.
+ no_object_weight (`float`, *optional*, defaults to 0.1):
+ Loss weight for the 'no object' class in panoptic/instance segmentation.
+ class_weight (`float`, *optional*, defaults to 2.0):
+ Loss weight for classification targets.
+ mask_weight (`float`, *optional*, defaults to 5.0):
+ Loss weight for mask prediction.
+ train_num_points (`int`, *optional*, defaults to 12544):
+ Number of points to sample for mask loss computation during training.
+ oversample_ratio (`float`, *optional*, defaults to 3.0):
+ Oversampling ratio used in point sampling for mask training.
+ importance_sample_ratio (`float`, *optional*, defaults to 0.75):
+ Ratio of points to sample based on importance during training.
+ num_queries (`int`, *optional*, defaults to 200):
+ Number of object queries in the Transformer.
+ num_register_tokens (`int`, *optional*, defaults to 4):
+ Number of learnable register tokens added to the transformer input.
+
+ Example:
+
+ ```python
+ >>> from transformers import EomtConfig, EomtForUniversalSegmentation
+
+ >>> # Initialize configuration
+ >>> config = EomtConfig()
+
+ >>> # Initialize model
+ >>> model = EomtForUniversalSegmentation(config)
+
+ >>> # Access config
+ >>> config = model.config
+ ```"""
+
+ model_type = "eomt"
+
+ hidden_size: int = 1024
+ num_hidden_layers: int = 24
+ num_attention_heads: int = 16
+ mlp_ratio: int = 4
+ hidden_act: str = "gelu"
+ hidden_dropout_prob: float | int = 0.0
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-6
+ image_size: int | list[int] | tuple[int, int] = 640
+ patch_size: int | list[int] | tuple[int, int] = 16
+ num_channels: int = 3
+ layerscale_value: float = 1.0
+ drop_path_rate: float | int = 0.0
+ num_upscale_blocks: int = 2
+ attention_dropout: float | int = 0.0
+ use_swiglu_ffn: bool = False
+ num_blocks: int = 4
+ no_object_weight: float = 0.1
+ class_weight: float = 2.0
+ mask_weight: float = 5.0
+ dice_weight: float = 5.0
+ train_num_points: int = 12544
+ oversample_ratio: float = 3.0
+ importance_sample_ratio: float = 0.75
+ num_queries: int = 200
+ num_register_tokens: int = 4
+
+ intermediate_size = AttributeError()
+ qkv_bias = AttributeError()
+ pooler_act = AttributeError()
+ pooler_output_size = AttributeError()
+ encoder_stride = AttributeError()
+ attention_probs_dropout_prob = AttributeError()
+
+ def __post_init__(self, **kwargs):
+ raise AttributeError("Not needed for Eomt")
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Class for outputs of [`EomtForUniversalSegmentationOutput`].
+
+ This output can be directly passed to [`~EomtImageProcessor.post_process_semantic_segmentation`] or
+ [`~EomtImageProcessor.post_process_instance_segmentation`] or
+ [`~EomtImageProcessor.post_process_panoptic_segmentation`] to compute final segmentation maps. Please, see
+ [`~EomtImageProcessor] for details regarding usage.
+ """
+)
+class EomtForUniversalSegmentationOutput(ModelOutput):
+ r"""
+ loss (`torch.Tensor`, *optional*):
+ The computed loss, returned when labels are present.
+ class_queries_logits (`torch.FloatTensor`):
+ A tensor of shape `(batch_size, num_queries, num_labels + 1)` representing the proposed classes for each
+ query. Note the `+ 1` is needed because we incorporate the null class.
+ masks_queries_logits (`torch.FloatTensor`):
+ A tensor of shape `(batch_size, num_queries, height, width)` representing the proposed masks for each
+ query.
+ last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Last hidden states (final feature map) of the last layer.
+ 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 + one for the output of each stage) of
+ shape `(batch_size, sequence_length, hidden_size)`. Hidden-states all layers of the model.
+ attentions (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+ Tuple of `tuple(torch.FloatTensor)` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+ sequence_length)`. Self and Cross Attentions weights from transformer decoder.
+ patch_offsets (`list[torch.Tensor]`, *optional*):
+ list of tuples indicating the image index and start and end positions of patches for semantic segmentation.
+ """
+
+ loss: torch.FloatTensor | None = None
+ class_queries_logits: torch.FloatTensor | None = None
+ masks_queries_logits: torch.FloatTensor | None = None
+ last_hidden_state: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ patch_offsets: list[torch.Tensor] | None = None
+
+
+class EomtLoss(Mask2FormerLoss):
+ pass
+
+
+class EomtPatchEmbeddings(Dinov2PatchEmbeddings):
+ pass
+
+
+class EomtEmbeddings(Dinov2Embeddings):
+ def __init__(self, config: EomtConfig) -> None:
+ nn.Module.__init__(self)
+
+ self.config = config
+ self.patch_size = config.patch_size
+
+ self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_size))
+ self.register_tokens = nn.Parameter(torch.zeros(1, config.num_register_tokens, config.hidden_size))
+
+ self.patch_embeddings = EomtPatchEmbeddings(config)
+ num_patches = self.patch_embeddings.num_patches
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ self.num_prefix_tokens = 1 + config.num_register_tokens # 1 for [CLS]
+ self.position_embeddings = nn.Embedding(num_patches, config.hidden_size)
+ self.register_buffer("position_ids", torch.arange(num_patches).expand((1, -1)), persistent=False)
+
+ def interpolate_pos_encoding(self):
+ raise AttributeError("Not needed for Eomt Model")
+
+ def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+ batch_size, _, _, _ = pixel_values.shape
+ target_dtype = self.patch_embeddings.projection.weight.dtype
+ embeddings = self.patch_embeddings(pixel_values.to(dtype=target_dtype))
+
+ cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+ register_tokens = self.register_tokens.expand(batch_size, -1, -1)
+
+ embeddings = embeddings + self.position_embeddings(self.position_ids)
+ embeddings = torch.cat([cls_tokens, register_tokens, embeddings], dim=1)
+
+ embeddings = self.dropout(embeddings)
+
+ return embeddings
+
+
+class EomtAttention(SiglipAttention):
+ pass
+
+
+class EomtLayerScale(Dinov2LayerScale):
+ pass
+
+
+class EomtLayer(Dinov2Layer):
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ hidden_states_norm = self.norm1(hidden_states)
+ self_attention_output, _ = self.attention(hidden_states_norm, attention_mask)
+ self_attention_output = self.layer_scale1(self_attention_output)
+
+ # first residual connection
+ hidden_states = self.drop_path(self_attention_output) + hidden_states
+
+ # in Eomt, layernorm is also applied after self-attention
+ layer_output = self.norm2(hidden_states)
+ layer_output = self.mlp(layer_output)
+ layer_output = self.layer_scale2(layer_output)
+
+ # second residual connection
+ layer_output = self.drop_path(layer_output) + hidden_states
+
+ return layer_output
+
+
+class EomtLayerNorm2d(nn.LayerNorm):
+ def __init__(self, num_channels, eps=1e-6, affine=True):
+ super().__init__(num_channels, eps=eps, elementwise_affine=affine)
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ hidden_state = hidden_state.permute(0, 2, 3, 1)
+ hidden_state = F.layer_norm(hidden_state, self.normalized_shape, self.weight, self.bias, self.eps)
+ hidden_state = hidden_state.permute(0, 3, 1, 2)
+ return hidden_state
+
+
+class EomtScaleLayer(nn.Module):
+ def __init__(self, config: EomtConfig):
+ super().__init__()
+ hidden_size = config.hidden_size
+ self.conv1 = nn.ConvTranspose2d(hidden_size, hidden_size, kernel_size=2, stride=2)
+ self.activation = ACT2FN[config.hidden_act]
+ self.conv2 = nn.Conv2d(
+ hidden_size,
+ hidden_size,
+ kernel_size=3,
+ padding=1,
+ groups=hidden_size,
+ bias=False,
+ )
+
+ self.layernorm2d = EomtLayerNorm2d(hidden_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.conv1(hidden_states)
+ hidden_states = self.activation(hidden_states)
+ hidden_states = self.conv2(hidden_states)
+ hidden_states = self.layernorm2d(hidden_states)
+ return hidden_states
+
+
+class EomtScaleBlock(nn.Module):
+ def __init__(self, config: EomtConfig):
+ super().__init__()
+ self.num_blocks = config.num_upscale_blocks
+ self.block = nn.ModuleList([EomtScaleLayer(config) for _ in range(self.num_blocks)])
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ for block in self.block:
+ hidden_states = block(hidden_states)
+ return hidden_states
+
+
+class EomtMaskHead(nn.Module):
+ def __init__(self, config: EomtConfig):
+ super().__init__()
+
+ hidden_size = config.hidden_size
+ self.fc1 = nn.Linear(hidden_size, hidden_size)
+ self.fc2 = nn.Linear(hidden_size, hidden_size)
+ self.fc3 = nn.Linear(hidden_size, hidden_size)
+ self.activation = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.activation(self.fc1(hidden_states))
+ hidden_states = self.activation(self.fc2(hidden_states))
+ hidden_states = self.fc3(hidden_states)
+ return hidden_states
+
+
+@auto_docstring
+class EomtPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config: EomtConfig
+ base_model_prefix = "eomt"
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+ supports_gradient_checkpointing = False
+ _no_split_modules = ["EomtLayer"]
+ _supports_sdpa = True
+ _can_record_outputs = {
+ "hidden_states": EomtLayer,
+ "attentions": EomtAttention,
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module: nn.Module) -> None:
+ std = self.config.initializer_range
+ if isinstance(module, (nn.Linear, nn.Conv2d, nn.ConvTranspose2d)):
+ init.kaiming_uniform_(module.weight, a=math.sqrt(5))
+ if module.bias is not None:
+ fan_in, _ = torch.nn.init._calculate_fan_in_and_fan_out(module.weight)
+ bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0
+ init.uniform_(module.bias, -bound, bound)
+ elif isinstance(module, nn.LayerNorm):
+ init.ones_(module.weight)
+ init.zeros_(module.bias)
+ elif isinstance(module, nn.Embedding):
+ init.normal_(module.weight, mean=0.0, std=1)
+ # Here we need the check explicitly, as we slice the weight in the `zeros_` call, so it looses the flag
+ if module.padding_idx is not None and not getattr(module.weight, "_is_hf_initialized", False):
+ init.zeros_(module.weight[module.padding_idx])
+ elif isinstance(module, EomtLayerScale):
+ if hasattr(module, "lambda1"):
+ init.constant_(module.lambda1, self.config.layerscale_value)
+ elif isinstance(module, EomtEmbeddings):
+ init.trunc_normal_(module.cls_token, mean=0.0, std=std)
+ init.zeros_(module.register_tokens)
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+ elif isinstance(module, EomtLoss):
+ empty_weight = torch.ones(module.num_labels + 1)
+ empty_weight[-1] = module.eos_coef
+ init.copy_(module.empty_weight, empty_weight)
+ elif isinstance(module, EomtForUniversalSegmentation):
+ init.ones_(module.attn_mask_probs)
+
+
+@auto_docstring(
+ custom_intro="""
+ The EoMT Model with head on top for instance/semantic/panoptic segmentation.
+ """
+)
+class EomtForUniversalSegmentation(Mask2FormerForUniversalSegmentation):
+ def __init__(self, config: EomtConfig):
+ PreTrainedModel.__init__(self, config)
+ self.config = config
+ self.num_hidden_layers = config.num_hidden_layers
+ self.embeddings = EomtEmbeddings(config)
+ self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+ self.query = nn.Embedding(config.num_queries, config.hidden_size)
+ self.layers = nn.ModuleList([EomtLayer(config) for _ in range(config.num_hidden_layers)])
+
+ self.upscale_block = EomtScaleBlock(config)
+ self.mask_head = EomtMaskHead(config)
+
+ self.class_predictor = nn.Linear(config.hidden_size, config.num_labels + 1)
+
+ self.grid_size = (config.image_size // config.patch_size, config.image_size // config.patch_size)
+ self.weight_dict: dict[str, float] = {
+ "loss_cross_entropy": config.class_weight,
+ "loss_mask": config.mask_weight,
+ "loss_dice": config.dice_weight,
+ }
+
+ self.criterion = EomtLoss(config=config, weight_dict=self.weight_dict)
+
+ self.register_buffer("attn_mask_probs", torch.ones(config.num_blocks))
+
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embeddings.patch_embeddings
+
+ def get_auxiliary_logits(self):
+ raise AttributeError("Note needed for Eomt Model.")
+
+ def predict(self, logits: torch.Tensor):
+ query_tokens = logits[:, : self.config.num_queries, :]
+ class_logits = self.class_predictor(query_tokens)
+
+ prefix_tokens = logits[:, self.config.num_queries + self.embeddings.num_prefix_tokens :, :]
+ prefix_tokens = prefix_tokens.transpose(1, 2)
+
+ prefix_tokens = prefix_tokens.reshape(prefix_tokens.shape[0], -1, *self.grid_size)
+
+ query_tokens = self.mask_head(query_tokens)
+ prefix_tokens = self.upscale_block(prefix_tokens)
+
+ mask_logits = torch.einsum("bqc, bchw -> bqhw", query_tokens, prefix_tokens)
+
+ return mask_logits, class_logits
+
+ @staticmethod
+ def _disable_attention_mask(attn_mask, prob, num_query_tokens, encoder_start_tokens, device):
+ if prob < 1:
+ # Generate random queries to disable based on the probs
+ random_queries = torch.rand(attn_mask.shape[0], num_query_tokens, device=device) > prob
+
+ # Disable attention to the query tokens, considering the prefix tokens
+ attn_mask[:, :num_query_tokens, encoder_start_tokens:][random_queries] = 1
+
+ return attn_mask
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: Tensor,
+ mask_labels: list[Tensor] | None = None,
+ class_labels: list[Tensor] | None = None,
+ patch_offsets: list[Tensor] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> EomtForUniversalSegmentationOutput:
+ r"""
+ mask_labels (`list[torch.Tensor]`, *optional*):
+ list of mask labels of shape `(num_labels, height, width)` to be fed to a model
+ class_labels (`list[torch.LongTensor]`, *optional*):
+ list of target class labels of shape `(num_labels, height, width)` to be fed to a model. They identify the
+ labels of `mask_labels`, e.g. the label of `mask_labels[i][j]` if `class_labels[i][j]`.
+ patch_offsets (`list[torch.Tensor]`, *optional*):
+ list of tuples indicating the image index and start and end positions of patches for semantic segmentation.
+ """
+
+ masks_queries_logits_per_layer, class_queries_logits_per_layer = (), ()
+ attention_mask = None
+
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ hidden_states = self.embeddings(pixel_values)
+
+ for idx, layer_module in enumerate(self.layers):
+ if idx == self.num_hidden_layers - self.config.num_blocks:
+ query = self.query.weight[None, :, :].expand(hidden_states.shape[0], -1, -1).to(hidden_states.device)
+ hidden_states = torch.cat((query, hidden_states), dim=1)
+
+ if idx >= self.num_hidden_layers - self.config.num_blocks and (
+ self.training or self.attn_mask_probs[idx - self.num_hidden_layers + self.config.num_blocks] > 0
+ ):
+ norm_hidden_states = self.layernorm(hidden_states)
+ masks_queries_logits, class_queries_logits = self.predict(norm_hidden_states)
+
+ masks_queries_logits_per_layer += (masks_queries_logits,)
+ class_queries_logits_per_layer += (class_queries_logits,)
+
+ attention_mask = torch.ones(
+ hidden_states.shape[0],
+ hidden_states.shape[1],
+ hidden_states.shape[1],
+ device=hidden_states.device,
+ dtype=torch.bool,
+ )
+
+ interpolated_logits = F.interpolate(masks_queries_logits, size=self.grid_size, mode="bilinear")
+ interpolated_logits = interpolated_logits.view(
+ interpolated_logits.size(0), interpolated_logits.size(1), -1
+ )
+
+ num_query_tokens = self.config.num_queries
+ encoder_start_tokens = num_query_tokens + self.embeddings.num_prefix_tokens
+
+ # Set attention mask for queries to focus on encoder tokens based on interpolated logits
+ attention_mask[:, :num_query_tokens, encoder_start_tokens:] = interpolated_logits > 0
+
+ # Disable attention mask for random query tokens.
+ attention_mask = self._disable_attention_mask(
+ attention_mask,
+ prob=self.attn_mask_probs[idx - self.num_hidden_layers + self.config.num_blocks],
+ num_query_tokens=num_query_tokens,
+ encoder_start_tokens=encoder_start_tokens,
+ device=attention_mask.device,
+ )
+
+ # Expand attention mask to 4d mask.
+ attention_mask = attention_mask[:, None, ...].expand(-1, self.config.num_attention_heads, -1, -1)
+ attention_mask = attention_mask.float().masked_fill(~attention_mask, -1e9)
+
+ hidden_states = layer_module(hidden_states, attention_mask)
+
+ sequence_output = self.layernorm(hidden_states)
+
+ masks_queries_logits, class_queries_logits = self.predict(sequence_output)
+ masks_queries_logits_per_layer += (masks_queries_logits,)
+ class_queries_logits_per_layer += (class_queries_logits,)
+
+ loss = None
+ if mask_labels is not None and class_labels is not None:
+ loss = 0.0
+ for masks_queries_logits, class_queries_logits in zip(
+ masks_queries_logits_per_layer, class_queries_logits_per_layer
+ ):
+ loss_dict = self.get_loss_dict(
+ masks_queries_logits=masks_queries_logits,
+ class_queries_logits=class_queries_logits,
+ mask_labels=mask_labels,
+ class_labels=class_labels,
+ auxiliary_predictions=None,
+ )
+ loss += self.get_loss(loss_dict)
+
+ return EomtForUniversalSegmentationOutput(
+ loss=loss,
+ masks_queries_logits=masks_queries_logits,
+ class_queries_logits=class_queries_logits,
+ last_hidden_state=sequence_output,
+ patch_offsets=patch_offsets,
+ )
+
+
+__all__ = ["EomtConfig", "EomtPreTrainedModel", "EomtForUniversalSegmentation"]
diff --git a/third_party/transformers/src/transformers/models/eurobert/__init__.py b/third_party/transformers/src/transformers/models/eurobert/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0212ce1dab27973c91b4feeccca8f12f0e24dbfd
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/eurobert/__init__.py
@@ -0,0 +1,27 @@
+# Copyright 2025 Nicolas Boizard, Duarte M. Alves, Hippolyte Gisserot-Boukhlef and the EuroBert team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_eurobert import *
+ from .modeling_eurobert import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/eurobert/configuration_eurobert.py b/third_party/transformers/src/transformers/models/eurobert/configuration_eurobert.py
new file mode 100644
index 0000000000000000000000000000000000000000..3c1585b7290ecc88618482e248b6af208cc9ac3c
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/eurobert/configuration_eurobert.py
@@ -0,0 +1,82 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/eurobert/modular_eurobert.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_eurobert.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 Nicolas Boizard, Duarte M. Alves, Hippolyte Gisserot-Boukhlef and the EuroBert team. All rights reserved.
+#
+#
+# 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.
+
+
+from ...configuration_utils import strict
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+from ..llama import LlamaConfig
+
+
+@auto_docstring(checkpoint="EuroBERT/EuroBERT-210m")
+@strict
+class EuroBertConfig(LlamaConfig):
+ r"""
+ mask_token_id (`int`, *optional*, defaults to 128002):
+ Mask token id.
+ classifier_pooling (`str`, *optional*, defaults to `"late"`):
+ The pooling strategy to use for the classifier. Can be one of ['bos', 'mean', 'late'].
+
+ ```python
+ >>> from transformers import EuroBertModel, EuroBertConfig
+
+ >>> # Initializing a EuroBert eurobert-base style configuration
+ >>> configuration = EuroBertConfig()
+
+ >>> # Initializing a model from the eurobert-base style configuration
+ >>> model = EuroBertModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "eurobert"
+
+ vocab_size: int = 128256
+ hidden_size: int = 768
+ intermediate_size: int = 3072
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 12
+ num_key_value_heads: int | None = None
+ hidden_act: str = "silu"
+ max_position_embeddings: int = 8192
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-05
+ bos_token_id: int | None = 128000
+ eos_token_id: int | list[int] | None = 128001
+ pad_token_id: int | None = 128001
+ mask_token_id: int = 128002
+ pretraining_tp: int = 1
+ tie_word_embeddings: bool = False
+ rope_parameters: RopeParameters | dict | None = None
+ attention_bias: bool = False
+ attention_dropout: int | float = 0.0
+ mlp_bias: bool = False
+ head_dim: int | None = None
+ classifier_pooling: str = "late"
+
+ def __post_init__(self, **kwargs):
+ if self.num_key_value_heads is None:
+ self.num_key_value_heads = self.num_attention_heads
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["EuroBertConfig"]
diff --git a/third_party/transformers/src/transformers/models/eurobert/modeling_eurobert.py b/third_party/transformers/src/transformers/models/eurobert/modeling_eurobert.py
new file mode 100644
index 0000000000000000000000000000000000000000..b93dd0649f14cb17c80595bb3aa105b5731f547d
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/eurobert/modeling_eurobert.py
@@ -0,0 +1,628 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/eurobert/modular_eurobert.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_eurobert.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 Nicolas Boizard, Duarte M. Alves, Hippolyte Gisserot-Boukhlef and the EuroBert team. All rights reserved.
+#
+#
+# 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.
+
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
+from ...masking_utils import create_bidirectional_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutput, MaskedLMOutput, SequenceClassifierOutput, TokenClassifierOutput
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import auto_docstring
+from ...utils.generic import TransformersKwargs, can_return_tuple, maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_eurobert import EuroBertConfig
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class EuroBertRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-5) -> None:
+ """
+ EuroBertRMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+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)
+
+
+@use_kernel_func_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+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)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class EuroBertAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: EuroBertConfig, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = False
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class EuroBertMLP(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=config.mlp_bias)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, x):
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+class EuroBertDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: EuroBertConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = EuroBertAttention(config=config, layer_idx=layer_idx)
+
+ self.mlp = EuroBertMLP(config)
+ self.input_layernorm = EuroBertRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = EuroBertRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ 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
+ return hidden_states
+
+
+@auto_docstring
+class EuroBertPreTrainedModel(PreTrainedModel):
+ config: EuroBertConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["EuroBertDecoderLayer"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": EuroBertDecoderLayer,
+ "attentions": EuroBertAttention,
+ }
+
+
+class EuroBertRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: EuroBertConfig, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: EuroBertConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@auto_docstring
+class EuroBertModel(EuroBertPreTrainedModel):
+ def __init__(self, config: EuroBertConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList(
+ [EuroBertDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = EuroBertRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = EuroBertRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutput:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
+
+ if position_ids is None:
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0)
+
+ bidirectional_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for encoder_layer in self.layers[: self.config.num_hidden_layers]:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask=bidirectional_mask,
+ position_embeddings=position_embeddings,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+@auto_docstring
+class EuroBertForMaskedLM(EuroBertPreTrainedModel):
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+ _tp_plan = {"lm_head": "colwise_gather_output"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+ def __init__(self, config: EuroBertConfig):
+ super().__init__(config)
+ self.model = EuroBertModel(config)
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, config.mlp_bias)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | MaskedLMOutput:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, EuroBertForMaskedLM
+
+ >>> model = EuroBertForMaskedLM.from_pretrained("EuroBERT/EuroBERT-210m")
+ >>> tokenizer = AutoTokenizer.from_pretrained("EuroBERT/EuroBERT-210m")
+
+ >>> text = "The capital of France is <|mask|>."
+ >>> inputs = tokenizer(text, return_tensors="pt")
+ >>> outputs = model(**inputs)
+
+ >>> # To get predictions for the mask:
+ >>> masked_index = inputs["input_ids"][0].tolist().index(tokenizer.mask_token_id)
+ >>> predicted_token_id = outputs.logits[0, masked_index].argmax(axis=-1)
+ >>> predicted_token = tokenizer.decode(predicted_token_id)
+ >>> print("Predicted token:", predicted_token)
+ Predicted token: Paris
+ ```"""
+ outputs: BaseModelOutput = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ logits = self.lm_head(outputs.last_hidden_state)
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return MaskedLMOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring
+class EuroBertForSequenceClassification(EuroBertPreTrainedModel):
+ def __init__(self, config: EuroBertConfig):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.classifier_pooling = config.classifier_pooling
+
+ self.model = EuroBertModel(config)
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.activation = nn.GELU()
+ self.classifier = nn.Linear(config.hidden_size, self.num_labels)
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | SequenceClassifierOutput:
+ encoder_output = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+ last_hidden_state = encoder_output[0]
+
+ if self.classifier_pooling in ["bos", "mean"]:
+ if self.classifier_pooling == "bos":
+ pooled_output = last_hidden_state[:, 0]
+
+ elif self.classifier_pooling == "mean":
+ if attention_mask is None:
+ pooled_output = last_hidden_state.mean(dim=1)
+ else:
+ attention_mask = attention_mask.to(last_hidden_state.device)
+ pooled_output = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(dim=1)
+ pooled_output /= attention_mask.sum(dim=1, keepdim=True)
+
+ pooled_output = self.dense(pooled_output)
+ pooled_output = self.activation(pooled_output)
+ logits = self.classifier(pooled_output)
+
+ elif self.classifier_pooling == "late":
+ x = self.dense(last_hidden_state)
+ x = self.activation(x)
+ logits = self.classifier(x)
+ if attention_mask is None:
+ logits = logits.mean(dim=1)
+ else:
+ attention_mask = attention_mask.to(logits.device)
+ logits = (logits * attention_mask.unsqueeze(-1)).sum(dim=1)
+ logits /= attention_mask.sum(dim=1, keepdim=True)
+
+ loss = None
+ if labels is not None:
+ labels = labels.to(logits.device)
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(logits, labels)
+
+ return SequenceClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=encoder_output.hidden_states,
+ attentions=encoder_output.attentions,
+ )
+
+
+@auto_docstring
+class EuroBertForTokenClassification(EuroBertPreTrainedModel):
+ def __init__(self, config: EuroBertConfig):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.model = EuroBertModel(config)
+
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.model.embed_tokens = value
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | TokenClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ outputs = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+ sequence_output = outputs[0]
+ logits = self.classifier(sequence_output)
+
+ loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+ return TokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "EuroBertPreTrainedModel",
+ "EuroBertModel",
+ "EuroBertForMaskedLM",
+ "EuroBertForSequenceClassification",
+ "EuroBertForTokenClassification",
+]
diff --git a/third_party/transformers/src/transformers/models/eurobert/modular_eurobert.py b/third_party/transformers/src/transformers/models/eurobert/modular_eurobert.py
new file mode 100644
index 0000000000000000000000000000000000000000..f0a2f1b0b479c4fa0103ac3387594d0349b0505f
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/eurobert/modular_eurobert.py
@@ -0,0 +1,362 @@
+# Copyright 2025 Nicolas Boizard, Duarte M. Alves, Hippolyte Gisserot-Boukhlef and the EuroBert team. All rights reserved.
+#
+#
+# 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.
+
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ...configuration_utils import strict
+from ...masking_utils import create_bidirectional_mask
+from ...modeling_outputs import BaseModelOutput, MaskedLMOutput, SequenceClassifierOutput, TokenClassifierOutput
+from ...modeling_rope_utils import RopeParameters
+from ...processing_utils import Unpack
+from ...utils import auto_docstring
+from ...utils.generic import TransformersKwargs, can_return_tuple
+from ..llama import LlamaConfig
+from ..llama.modeling_llama import LlamaAttention, LlamaModel, LlamaPreTrainedModel, LlamaRMSNorm
+
+
+@auto_docstring(checkpoint="EuroBERT/EuroBERT-210m")
+@strict
+class EuroBertConfig(LlamaConfig):
+ r"""
+ mask_token_id (`int`, *optional*, defaults to 128002):
+ Mask token id.
+ classifier_pooling (`str`, *optional*, defaults to `"late"`):
+ The pooling strategy to use for the classifier. Can be one of ['bos', 'mean', 'late'].
+
+ ```python
+ >>> from transformers import EuroBertModel, EuroBertConfig
+
+ >>> # Initializing a EuroBert eurobert-base style configuration
+ >>> configuration = EuroBertConfig()
+
+ >>> # Initializing a model from the eurobert-base style configuration
+ >>> model = EuroBertModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "eurobert"
+
+ vocab_size: int = 128256
+ hidden_size: int = 768
+ intermediate_size: int = 3072
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 12
+ num_key_value_heads: int | None = None
+ hidden_act: str = "silu"
+ max_position_embeddings: int = 8192
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-05
+ bos_token_id: int | None = 128000
+ eos_token_id: int | list[int] | None = 128001
+ pad_token_id: int | None = 128001
+ mask_token_id: int = 128002
+ pretraining_tp: int = 1
+ tie_word_embeddings: bool = False
+ rope_parameters: RopeParameters | dict | None = None
+ attention_bias: bool = False
+ attention_dropout: int | float = 0.0
+ mlp_bias: bool = False
+ head_dim: int | None = None
+ classifier_pooling: str = "late"
+
+ def __post_init__(self, **kwargs):
+ if self.num_key_value_heads is None:
+ self.num_key_value_heads = self.num_attention_heads
+ super().__post_init__(**kwargs)
+
+
+class EuroBertRMSNorm(LlamaRMSNorm):
+ def __init__(self, hidden_size, eps=1e-5):
+ super().__init__(hidden_size, eps)
+
+
+class EuroBertAttention(LlamaAttention):
+ def __init__(self, config: EuroBertConfig, layer_idx: int):
+ super().__init__(config, layer_idx)
+ self.is_causal = False
+
+
+class EuroBertPreTrainedModel(LlamaPreTrainedModel):
+ pass
+
+
+class EuroBertModel(LlamaModel):
+ def forward(
+ self,
+ input_ids: torch.LongTensor = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutput:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
+
+ if position_ids is None:
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0)
+
+ bidirectional_mask = create_bidirectional_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for encoder_layer in self.layers[: self.config.num_hidden_layers]:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask=bidirectional_mask,
+ position_embeddings=position_embeddings,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+@auto_docstring
+class EuroBertForMaskedLM(EuroBertPreTrainedModel):
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+ _tp_plan = {"lm_head": "colwise_gather_output"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+ def __init__(self, config: EuroBertConfig):
+ super().__init__(config)
+ self.model = EuroBertModel(config)
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, config.mlp_bias)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | MaskedLMOutput:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, EuroBertForMaskedLM
+
+ >>> model = EuroBertForMaskedLM.from_pretrained("EuroBERT/EuroBERT-210m")
+ >>> tokenizer = AutoTokenizer.from_pretrained("EuroBERT/EuroBERT-210m")
+
+ >>> text = "The capital of France is <|mask|>."
+ >>> inputs = tokenizer(text, return_tensors="pt")
+ >>> outputs = model(**inputs)
+
+ >>> # To get predictions for the mask:
+ >>> masked_index = inputs["input_ids"][0].tolist().index(tokenizer.mask_token_id)
+ >>> predicted_token_id = outputs.logits[0, masked_index].argmax(axis=-1)
+ >>> predicted_token = tokenizer.decode(predicted_token_id)
+ >>> print("Predicted token:", predicted_token)
+ Predicted token: Paris
+ ```"""
+ outputs: BaseModelOutput = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ logits = self.lm_head(outputs.last_hidden_state)
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return MaskedLMOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring
+class EuroBertForSequenceClassification(EuroBertPreTrainedModel):
+ def __init__(self, config: EuroBertConfig):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.classifier_pooling = config.classifier_pooling
+
+ self.model = EuroBertModel(config)
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.activation = nn.GELU()
+ self.classifier = nn.Linear(config.hidden_size, self.num_labels)
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor] | SequenceClassifierOutput:
+ encoder_output = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+ last_hidden_state = encoder_output[0]
+
+ if self.classifier_pooling in ["bos", "mean"]:
+ if self.classifier_pooling == "bos":
+ pooled_output = last_hidden_state[:, 0]
+
+ elif self.classifier_pooling == "mean":
+ if attention_mask is None:
+ pooled_output = last_hidden_state.mean(dim=1)
+ else:
+ attention_mask = attention_mask.to(last_hidden_state.device)
+ pooled_output = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(dim=1)
+ pooled_output /= attention_mask.sum(dim=1, keepdim=True)
+
+ pooled_output = self.dense(pooled_output)
+ pooled_output = self.activation(pooled_output)
+ logits = self.classifier(pooled_output)
+
+ elif self.classifier_pooling == "late":
+ x = self.dense(last_hidden_state)
+ x = self.activation(x)
+ logits = self.classifier(x)
+ if attention_mask is None:
+ logits = logits.mean(dim=1)
+ else:
+ attention_mask = attention_mask.to(logits.device)
+ logits = (logits * attention_mask.unsqueeze(-1)).sum(dim=1)
+ logits /= attention_mask.sum(dim=1, keepdim=True)
+
+ loss = None
+ if labels is not None:
+ labels = labels.to(logits.device)
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(logits, labels)
+
+ return SequenceClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=encoder_output.hidden_states,
+ attentions=encoder_output.attentions,
+ )
+
+
+@auto_docstring
+class EuroBertForTokenClassification(EuroBertPreTrainedModel):
+ def __init__(self, config: EuroBertConfig):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.model = EuroBertModel(config)
+
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.embed_tokens
+
+ def set_input_embeddings(self, value):
+ self.model.embed_tokens = value
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | TokenClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ outputs = self.model(
+ input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+ sequence_output = outputs[0]
+ logits = self.classifier(sequence_output)
+
+ loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+ return TokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "EuroBertConfig",
+ "EuroBertPreTrainedModel",
+ "EuroBertModel",
+ "EuroBertForMaskedLM",
+ "EuroBertForSequenceClassification",
+ "EuroBertForTokenClassification",
+]
diff --git a/third_party/transformers/src/transformers/models/falcon/__init__.py b/third_party/transformers/src/transformers/models/falcon/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f9789767f11402264660b5dec0b5cae2466ee9d8
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/falcon/__init__.py
@@ -0,0 +1,27 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_falcon import *
+ from .modeling_falcon import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/falcon/configuration_falcon.py b/third_party/transformers/src/transformers/models/falcon/configuration_falcon.py
new file mode 100644
index 0000000000000000000000000000000000000000..5ff221e809498cc9036b18beb27e16f164216044
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/falcon/configuration_falcon.py
@@ -0,0 +1,110 @@
+# Copyright 2023 the Falcon authors and HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Falcon configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="tiiuae/falcon-7b")
+@strict
+class FalconConfig(PreTrainedConfig):
+ r"""
+ num_ln_in_parallel_attn (`int`, *optional*):
+ Set to 2 if separate layer norms are to be used for the MLP and the attention output when using parallel
+ attention, otherwise, 1.
+ alibi (`bool`, *optional*, defaults to `False`):
+ Whether to use ALiBi positional biases during self-attention.
+ new_decoder_architecture (`bool`, *optional*, defaults to `False`):
+ Whether to use the new (Falcon-40B) decoder architecture. If `True`, the `multi_query` and `parallel_attn`
+ arguments are ignored, as the new decoder always uses parallel attention.
+ multi_query (`bool`, *optional*, defaults to `True`):
+ Whether to use multi-query attention in the decoder. Ignored when `new_decoder_architecture` is `True`.
+ parallel_attn (`bool`, *optional*, defaults to `True`):
+ Whether to compute attention in parallel with the feedforward layer. If False, they are consecutive
+ instead, as in the original Transformer architecture. Ignored when `new_decoder_architecture` is `True`.
+ bias (`bool`, *optional*, defaults to `False`):
+ Whether to use bias on Linear layers.
+ ffn_hidden_size (`int`, *optional*):
+ The hidden size of the feedforward layer in the Transformer decoder.
+ defaults to 4x hidden dim
+ activation (`str`, *optional*, defaults to `"gelu"`):
+ The activation function used in the feedforward layer.
+
+ Example:
+
+ ```python
+ >>> from transformers import FalconModel, FalconConfig
+
+ >>> # Initializing a small (2-layer) Falcon configuration
+ >>> configuration = FalconConfig(num_hidden_layers=2)
+
+ >>> # Initializing a model from the small configuration
+ >>> model = FalconModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "falcon"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ vocab_size: int = 65024
+ hidden_size: int = 4544
+ num_hidden_layers: int = 32
+ num_attention_heads: int = 71
+ num_ln_in_parallel_attn: int | None = None
+ layer_norm_epsilon: float | None = 1e-5
+ initializer_range: float = 0.02
+ use_cache: bool = True
+ hidden_dropout: float | int | None = 0.0
+ attention_dropout: float | int | None = 0.0
+ num_kv_heads: int | None = None
+ alibi: bool | None = False
+ new_decoder_architecture: bool | None = False
+ multi_query: bool | None = True
+ parallel_attn: bool | None = True
+ bias: bool | None = False
+ max_position_embeddings: int = 2048
+ rope_parameters: RopeParameters | dict | None = None
+ bos_token_id: int | None = 11
+ eos_token_id: int | list[int] | None = 11
+ pad_token_id: int | None = None
+ ffn_hidden_size: int | None = None
+ activation: str | None = "gelu"
+ tie_word_embeddings: bool = True
+
+ def __post_init__(self, **kwargs):
+ # Backward compatibility with n_embed kwarg
+ n_embed = kwargs.pop("n_embed", None)
+ self.hidden_size = self.hidden_size if n_embed is None else n_embed
+ self.num_kv_heads = self.num_attention_heads if self.num_kv_heads is None else self.num_kv_heads
+ if self.ffn_hidden_size is None:
+ self.ffn_hidden_size = self.hidden_size * 4
+
+ super().__post_init__(**kwargs)
+
+ @property
+ def head_dim(self):
+ return self.hidden_size // self.num_attention_heads
+
+ @property
+ def rotary(self):
+ return not self.alibi
+
+
+__all__ = ["FalconConfig"]
diff --git a/third_party/transformers/src/transformers/models/falcon/convert_custom_code_checkpoint.py b/third_party/transformers/src/transformers/models/falcon/convert_custom_code_checkpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..0da817c3ffa73907c0215be12377f08fb5729a85
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/falcon/convert_custom_code_checkpoint.py
@@ -0,0 +1,74 @@
+import json
+from argparse import ArgumentParser
+from pathlib import Path
+
+
+"""
+This script converts Falcon custom code checkpoints to modern Falcon checkpoints that use code in the Transformers
+library. After conversion, performance (especially for generation) should improve and the checkpoint can be loaded
+without needing trust_remote_code=True.
+"""
+
+if __name__ == "__main__":
+ parser = ArgumentParser()
+ parser.add_argument(
+ "--checkpoint_dir",
+ type=Path,
+ required=True,
+ help="Directory containing a custom code checkpoint to convert to a modern Falcon checkpoint.",
+ )
+ args = parser.parse_args()
+
+ if not args.checkpoint_dir.is_dir():
+ raise ValueError("--checkpoint_dir argument should be a directory!")
+
+ if (
+ not (args.checkpoint_dir / "configuration_RW.py").is_file()
+ or not (args.checkpoint_dir / "modelling_RW.py").is_file()
+ ):
+ raise ValueError(
+ "The model directory should contain configuration_RW.py and modelling_RW.py files! Are you sure this is a custom code checkpoint?"
+ )
+ (args.checkpoint_dir / "configuration_RW.py").unlink()
+ (args.checkpoint_dir / "modelling_RW.py").unlink()
+
+ config = args.checkpoint_dir / "config.json"
+ text = config.read_text()
+ text = text.replace("RWForCausalLM", "FalconForCausalLM")
+ text = text.replace("RefinedWebModel", "falcon")
+ text = text.replace("RefinedWeb", "falcon")
+ json_config = json.loads(text)
+ del json_config["auto_map"]
+
+ if "n_head" in json_config:
+ json_config["num_attention_heads"] = json_config.pop("n_head")
+ if "n_layer" in json_config:
+ json_config["num_hidden_layers"] = json_config.pop("n_layer")
+ if "n_head_kv" in json_config:
+ json_config["num_kv_heads"] = json_config.pop("n_head_kv")
+ json_config["new_decoder_architecture"] = True
+ else:
+ json_config["new_decoder_architecture"] = False
+ bos_token_id = json_config.get("bos_token_id", 1)
+ eos_token_id = json_config.get("eos_token_id", 2)
+ config.unlink()
+ config.write_text(json.dumps(json_config, indent=2, sort_keys=True))
+
+ tokenizer_config = args.checkpoint_dir / "tokenizer_config.json"
+ if tokenizer_config.is_file():
+ text = tokenizer_config.read_text()
+ json_config = json.loads(text)
+ if json_config["tokenizer_class"] == "PreTrainedTokenizerFast":
+ json_config["model_input_names"] = ["input_ids", "attention_mask"]
+ tokenizer_config.unlink()
+ tokenizer_config.write_text(json.dumps(json_config, indent=2, sort_keys=True))
+
+ generation_config_path = args.checkpoint_dir / "generation_config.json"
+ generation_dict = {
+ "_from_model_config": True,
+ "bos_token_id": bos_token_id,
+ "eos_token_id": eos_token_id,
+ "transformers_version": "4.33.0.dev0",
+ }
+ generation_config_path.write_text(json.dumps(generation_dict, indent=2, sort_keys=True))
+ print("Done! Please double-check that the new checkpoint works as expected.")
diff --git a/third_party/transformers/src/transformers/models/falcon/modeling_falcon.py b/third_party/transformers/src/transformers/models/falcon/modeling_falcon.py
new file mode 100644
index 0000000000000000000000000000000000000000..016b3209b6b127089f7c126430205baa36747962
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/falcon/modeling_falcon.py
@@ -0,0 +1,1239 @@
+# Copyright 2023 the Falcon authors and HuggingFace Inc. team. All rights reserved.
+#
+# 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 Falcon model."""
+
+import math
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss
+from torch.nn import functional as F
+
+from ... import initialization as init
+from ...activations import get_activation
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import flash_attn_supports_top_left_mask, is_flash_attn_available
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutputWithPastAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions,
+ QuestionAnsweringModelOutput,
+ SequenceClassifierOutputWithPast,
+ TokenClassifierOutput,
+)
+from ...modeling_rope_utils import (
+ ROPE_INIT_FUNCTIONS,
+ dynamic_rope_update,
+)
+from ...modeling_utils import PreTrainedModel
+from ...utils import (
+ auto_docstring,
+ logging,
+)
+from ...utils.generic import maybe_autocast
+from .configuration_falcon import FalconConfig
+
+
+if is_flash_attn_available():
+ from ...modeling_flash_attention_utils import _flash_attention_forward
+
+logger = logging.get_logger(__name__)
+
+
+# NOTE(Hesslow): Unfortunately we did not fuse matmul and bias during training, this means that there's one additional quantization to bfloat16 between the operations.
+# In order not to degrade the quality of our HF-port, we keep these characteristics in the final model.
+class FalconLinear(nn.Linear):
+ def forward(self, input: torch.Tensor) -> torch.Tensor:
+ hidden_states = input @ self.weight.T
+ if self.bias is None:
+ return hidden_states
+ return hidden_states + self.bias
+
+
+# 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, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+# Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Falcon
+class FalconRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: FalconConfig, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: FalconConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+def build_alibi_tensor(attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor:
+ batch_size, seq_length = attention_mask.shape
+ closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
+ base = torch.tensor(
+ 2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32
+ )
+ powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32)
+ slopes = torch.pow(base, powers)
+
+ if closest_power_of_2 != num_heads:
+ extra_base = torch.tensor(
+ 2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32
+ )
+ num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
+ extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=attention_mask.device, dtype=torch.int32)
+ slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
+
+ # Note: alibi will added to the attention bias that will be applied to the query, key product of attention
+ # => therefore alibi will have to be of shape (batch_size, num_heads, query_length, key_length)
+ # => here we set (batch_size=1, num_heads=num_heads, query_length=1, key_length=max_length)
+ # => the query_length dimension will then be broadcasted correctly
+ # This is more or less identical to T5's relative position bias:
+ # https://github.com/huggingface/transformers/blob/f681437203baa7671de3174b0fa583c349d9d5e1/src/transformers/models/t5/modeling_t5.py#L527
+ arange_tensor = ((attention_mask.cumsum(dim=-1) - 1) * attention_mask)[:, None, :]
+ alibi = slopes[..., None].bfloat16() * arange_tensor
+ return alibi.reshape(batch_size * num_heads, 1, seq_length).to(dtype)
+
+
+# Copied from transformers.models.bloom.modeling_bloom.dropout_add
+def dropout_add(x: torch.Tensor, residual: torch.Tensor, prob: float, training: bool) -> torch.Tensor:
+ """
+ Dropout add function
+
+ Args:
+ x (`torch.tensor`):
+ input tensor
+ residual (`torch.tensor`):
+ residual tensor
+ prob (`float`):
+ dropout probability
+ training (`bool`):
+ training mode
+ """
+ out = F.dropout(x, p=prob, training=training)
+ out = residual + out
+ return out
+
+
+class FalconAttention(nn.Module):
+ def __init__(self, config: FalconConfig, layer_idx=None):
+ 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.split_size = self.hidden_size
+ self.hidden_dropout = config.hidden_dropout
+ self.max_position_embeddings = config.max_position_embeddings
+
+ self.is_causal = True
+ self.layer_idx = layer_idx
+ if layer_idx is None:
+ logger.warning_once(
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class."
+ )
+
+ 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} and `num_heads`:"
+ f" {self.num_heads})."
+ )
+
+ # Layer-wise attention scaling
+ self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim)
+ self.beta = self.inv_norm_factor
+ if config.new_decoder_architecture:
+ qkv_out_dim = (config.num_kv_heads * 2 + config.num_attention_heads) * self.head_dim
+ elif config.multi_query:
+ qkv_out_dim = self.hidden_size + 2 * self.head_dim
+ else:
+ qkv_out_dim = 3 * self.hidden_size
+ self.query_key_value = FalconLinear(self.hidden_size, qkv_out_dim, bias=config.bias)
+ self.new_decoder_architecture = config.new_decoder_architecture
+ self.multi_query = config.multi_query
+ self.dense = FalconLinear(self.hidden_size, self.hidden_size, bias=config.bias)
+ self.attention_dropout = nn.Dropout(config.attention_dropout)
+ self.num_kv_heads = config.num_kv_heads if (self.new_decoder_architecture or not self.multi_query) else 1
+
+ def _split_heads(self, fused_qkv: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ """
+ Split the last dimension into (num_heads, head_dim), results share same memory storage as `fused_qkv`
+
+ Args:
+ fused_qkv (`torch.tensor`): [batch_size, seq_length, num_heads * 3 * head_dim]
+
+ Returns:
+ query: [batch_size, seq_length, num_heads, head_dim] key: [batch_size, seq_length, num_heads, head_dim]
+ value: [batch_size, seq_length, num_heads, head_dim]
+ """
+ if self.new_decoder_architecture:
+ batch, seq_len, _ = fused_qkv.shape
+ qkv = fused_qkv.view(batch, seq_len, -1, self.num_heads // self.num_kv_heads + 2, self.head_dim)
+ query = qkv[:, :, :, :-2]
+ key = qkv[:, :, :, [-2]]
+ value = qkv[:, :, :, [-1]]
+ key = torch.broadcast_to(key, query.shape)
+ value = torch.broadcast_to(value, query.shape)
+
+ query, key, value = [x.flatten(2, 3) for x in (query, key, value)]
+ return query, key, value
+ elif not self.multi_query:
+ batch_size, seq_length, three_times_hidden_size = fused_qkv.shape
+ fused_qkv = fused_qkv.view(batch_size, seq_length, self.num_heads, 3, self.head_dim)
+ return fused_qkv[..., 0, :], fused_qkv[..., 1, :], fused_qkv[..., 2, :]
+ else:
+ batch_size, seq_length, three_times_hidden_size = fused_qkv.shape
+ fused_qkv = fused_qkv.view(batch_size, seq_length, self.num_heads + 2, self.head_dim)
+ return fused_qkv[..., :-2, :], fused_qkv[..., [-2], :], fused_qkv[..., [-1], :]
+
+ # Copied from transformers.models.bloom.modeling_bloom.BloomAttention._merge_heads
+ def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:
+ """
+ Merge heads together over the last dimension
+
+ Args:
+ x (`torch.tensor`): [batch_size * num_heads, seq_length, head_dim]
+
+ Returns:
+ torch.tensor: [batch_size, seq_length, num_heads * head_dim]
+ """
+ # What we want to achieve is:
+ # batch_size * num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads * head_dim
+ batch_size_and_num_heads, seq_length, _ = x.shape
+ batch_size = batch_size_and_num_heads // self.num_heads
+
+ # First view to decompose the batch size
+ # batch_size * num_heads, seq_length, head_dim -> batch_size, num_heads, seq_length, head_dim
+ x = x.view(batch_size, self.num_heads, seq_length, self.head_dim)
+
+ # batch_size, num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads, head_dim
+ x = x.permute(0, 2, 1, 3)
+
+ # batch_size, seq_length, num_heads, head_dim -> batch_size, seq_length, num_heads * head_dim
+ return x.reshape(batch_size, seq_length, self.num_heads * self.head_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ alibi: torch.Tensor | None,
+ attention_mask: torch.Tensor,
+ position_ids: torch.LongTensor | None = None,
+ layer_past: Cache | None = None,
+ use_cache: bool = False,
+ output_attentions: bool = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs,
+ ):
+ fused_qkv = self.query_key_value(hidden_states) # [batch_size, seq_length, 3 x hidden_size]
+ num_kv_heads = self.num_heads if self.new_decoder_architecture else self.num_kv_heads
+ # 3 x [batch_size, seq_length, num_heads, head_dim]
+ (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv)
+
+ batch_size, query_length, _, _ = query_layer.shape
+
+ query_layer = query_layer.transpose(1, 2).reshape(batch_size, self.num_heads, query_length, self.head_dim)
+ key_layer = key_layer.transpose(1, 2).reshape(batch_size, num_kv_heads, query_length, self.head_dim)
+ value_layer = value_layer.transpose(1, 2).reshape(batch_size, num_kv_heads, query_length, self.head_dim)
+
+ if alibi is None:
+ cos, sin = position_embeddings
+ query_layer, key_layer = apply_rotary_pos_emb(query_layer, key_layer, cos, sin)
+
+ if layer_past is not None:
+ key_layer, value_layer = layer_past.update(key_layer, value_layer, self.layer_idx)
+
+ kv_length = key_layer.shape[-2]
+
+ if alibi is None:
+ if self.config._attn_implementation == "sdpa" and not output_attentions:
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this if statement instead of an
+ # inline conditional assignment to support both torch.compile's `dynamic=True` and `fullgraph=True`
+ # The query_length > 1 is necessary to match with a bidirectional attention mask we do not have
+ # a causal pattern in those cases.
+ is_causal = self.is_causal and attention_mask is None and query_length > 1
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
+ query_layer,
+ key_layer,
+ value_layer,
+ attn_mask=attention_mask,
+ dropout_p=0.0,
+ is_causal=is_causal,
+ )
+ attention_scores = None
+ else:
+ attention_scores = query_layer @ key_layer.transpose(-1, -2)
+ attention_scores /= math.sqrt(self.head_dim)
+
+ attention_scores = F.softmax(attention_scores + attention_mask, dim=-1, dtype=hidden_states.dtype)
+ # It is unclear why dropout is not applied here (while it is with alibi).
+ attn_output = attention_scores @ value_layer
+
+ attn_output = attn_output.view(batch_size, self.num_heads, query_length, self.head_dim)
+ attn_output = attn_output.permute(0, 2, 1, 3)
+ attn_output = attn_output.reshape(batch_size, query_length, self.num_heads * self.head_dim)
+
+ attn_output = self.dense(attn_output)
+
+ return attn_output, attention_scores
+
+ else:
+ if self.config._attn_implementation == "sdpa" and not output_attentions:
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this if statement instead of an
+ # inline conditional assignment to support both torch.compile's `dynamic=True` and `fullgraph=True`
+ is_causal = self.is_causal and attention_mask is None and query_length > 1
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
+ query_layer,
+ key_layer,
+ value_layer,
+ attn_mask=attention_mask,
+ dropout_p=self.attention_dropout.p if self.training else 0.0,
+ is_causal=is_causal,
+ )
+ attention_probs = None
+ attn_output = attn_output.transpose(1, 2)
+ attn_output = attn_output.reshape(batch_size, query_length, self.num_heads * self.head_dim)
+
+ attn_output = self.dense(attn_output)
+ else:
+ matmul_result = query_layer @ key_layer.transpose(-1, -2)
+
+ # change view to [batch_size, num_heads, q_length, kv_length]
+ attention_scores = matmul_result.view(batch_size, self.num_heads, query_length, kv_length)
+
+ # cast attention scores to fp32, compute scaled softmax and cast back to initial dtype - [batch_size, num_heads, q_length, kv_length]
+ input_dtype = attention_scores.dtype
+ # `float16` has a minimum value of -65504.0, whereas `bfloat16` and `float32` have a minimum value of `-3.4e+38`
+ if input_dtype == torch.float16 or input_dtype == torch.bfloat16:
+ attention_scores = attention_scores.to(torch.float32)
+
+ attention_logits = attention_scores + alibi.view(batch_size, self.num_heads, 1, -1)
+ attention_logits *= self.inv_norm_factor
+ attention_probs = F.softmax(attention_logits + attention_mask, dim=-1, dtype=hidden_states.dtype)
+ # [batch_size, num_heads, q_length, kv_length]
+ attention_probs = self.attention_dropout(attention_probs)
+
+ # change view [batch_size, num_heads, q_length, kv_length]
+ attention_probs_reshaped = attention_probs.view(batch_size, self.num_heads, query_length, kv_length)
+
+ # matmul: [batch_size * num_heads, q_length, head_dim]
+ attn_output = (attention_probs_reshaped @ value_layer).flatten(0, 1)
+
+ # change view [batch_size, q_length, num_heads * head_dim]
+ attn_output = self._merge_heads(attn_output)
+
+ attn_output = self.dense(attn_output)
+
+ return attn_output, attention_probs
+
+
+class FalconFlashAttention2(FalconAttention):
+ """
+ Falcon flash attention module. This module inherits from `FalconAttention` 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 __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignment, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
+ # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
+ self._flash_attn_uses_top_left_mask = flash_attn_supports_top_left_mask()
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ alibi: torch.Tensor | None,
+ attention_mask: torch.Tensor,
+ position_ids: torch.LongTensor | None = None,
+ layer_past: Cache | None = None,
+ use_cache: bool = False,
+ output_attentions: bool = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs,
+ ):
+ fused_qkv = self.query_key_value(hidden_states) # [batch_size, seq_length, 3 x hidden_size]
+ num_kv_heads = self.num_heads if self.new_decoder_architecture else self.num_kv_heads
+ # 3 x [batch_size, seq_length, num_heads, head_dim]
+ (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv)
+
+ batch_size, query_length, _, _ = query_layer.shape
+
+ query_layer = query_layer.transpose(1, 2).reshape(batch_size, self.num_heads, query_length, self.head_dim)
+ key_layer = key_layer.transpose(1, 2).reshape(batch_size, num_kv_heads, query_length, self.head_dim)
+ value_layer = value_layer.transpose(1, 2).reshape(batch_size, num_kv_heads, query_length, self.head_dim)
+
+ if alibi is None:
+ cos, sin = position_embeddings
+ query_layer, key_layer = apply_rotary_pos_emb(query_layer, key_layer, cos, sin)
+
+ if layer_past is not None:
+ key_layer, value_layer = layer_past.update(key_layer, value_layer, self.layer_idx)
+
+ # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
+ # to be able to avoid many of these transpose/reshape/view.
+ query_layer = query_layer.transpose(1, 2)
+ key_layer = key_layer.transpose(1, 2)
+ value_layer = value_layer.transpose(1, 2)
+
+ if alibi is not None:
+ raise ValueError("`alibi` is not supported when `use_flash_attn` is True")
+
+ attn_dropout = self.config.attention_dropout if self.training else 0.0
+
+ # 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_layer.dtype
+ device_type = query_layer.device.type if query_layer.device.type != "mps" else "cpu"
+ if input_dtype == torch.float32:
+ if torch.is_autocast_enabled(device_type):
+ target_dtype = torch.get_autocast_dtype(device_type)
+ # Handle the case where the model is quantized
+ elif hasattr(self.config, "_is_quantized"):
+ target_dtype = self.config.dtype
+ else:
+ target_dtype = self.query_key_value.weight.dtype
+
+ logger.warning_once(
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
+ f" {target_dtype}."
+ )
+
+ query_layer = query_layer.to(target_dtype)
+ key_layer = key_layer.to(target_dtype)
+ value_layer = value_layer.to(target_dtype)
+
+ attn_output = _flash_attention_forward(
+ query_layer,
+ key_layer,
+ value_layer,
+ attention_mask,
+ query_length,
+ position_ids=position_ids,
+ dropout=attn_dropout,
+ is_causal=self.is_causal,
+ use_top_left_mask=self._flash_attn_uses_top_left_mask,
+ )
+
+ attn_weights = attn_output.reshape(batch_size, query_length, self.num_heads * self.head_dim)
+ attn_output = self.dense(attn_weights)
+
+ if not output_attentions:
+ attn_weights = None
+
+ return attn_output, attn_weights
+
+
+class FalconMLP(nn.Module):
+ def __init__(self, config: FalconConfig):
+ super().__init__()
+ hidden_size = config.hidden_size
+
+ self.dense_h_to_4h = FalconLinear(hidden_size, config.ffn_hidden_size, bias=config.bias)
+ self.act = get_activation(config.activation)
+ self.dense_4h_to_h = FalconLinear(config.ffn_hidden_size, hidden_size, bias=config.bias)
+ self.hidden_dropout = config.hidden_dropout
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = self.act(self.dense_h_to_4h(x))
+ x = self.dense_4h_to_h(x)
+ return x
+
+
+FALCON_ATTENTION_CLASSES = {
+ "eager": FalconAttention,
+ "sdpa": FalconAttention, # FalconAttention originally implemented both a forward with & without SDPA
+ "flash_attention_2": FalconFlashAttention2,
+}
+
+
+class FalconDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: FalconConfig, layer_idx=None):
+ super().__init__()
+ hidden_size = config.hidden_size
+ self.num_heads = config.num_attention_heads
+
+ self.self_attention = FALCON_ATTENTION_CLASSES[config._attn_implementation](config, layer_idx)
+ self.mlp = FalconMLP(config)
+ self.hidden_dropout = config.hidden_dropout
+ self.config = config
+
+ if config.num_ln_in_parallel_attn is None and config.new_decoder_architecture:
+ config.num_ln_in_parallel_attn = 2
+
+ if not config.parallel_attn:
+ self.post_attention_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+ self.input_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+ else:
+ if config.num_ln_in_parallel_attn == 2:
+ # The layer norm before self-attention
+ self.ln_attn = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+ # The layer norm before the MLP
+ self.ln_mlp = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+ else:
+ self.input_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ alibi: torch.Tensor | None,
+ attention_mask: torch.Tensor,
+ position_ids: torch.LongTensor | None = None,
+ layer_past: Cache | tuple[torch.Tensor, torch.Tensor] | None = None,
+ use_cache: bool = False,
+ output_attentions: bool = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs,
+ ):
+ residual = hidden_states
+
+ if self.config.new_decoder_architecture and self.config.num_ln_in_parallel_attn == 2:
+ attention_layernorm_out = self.ln_attn(hidden_states)
+ mlp_layernorm_out = self.ln_mlp(hidden_states)
+ else:
+ attention_layernorm_out = self.input_layernorm(hidden_states)
+
+ # Self attention.
+ attention_output, attn_weights = self.self_attention(
+ attention_layernorm_out,
+ layer_past=layer_past,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ alibi=alibi,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ position_embeddings=position_embeddings,
+ )
+
+ if not self.config.new_decoder_architecture:
+ if self.config.parallel_attn:
+ mlp_layernorm_out = attention_layernorm_out
+ else:
+ residual = dropout_add(
+ attention_output, residual, self.config.attention_dropout, training=self.training
+ )
+ mlp_layernorm_out = self.post_attention_layernorm(residual)
+
+ if (
+ self.config.new_decoder_architecture
+ and self.config.parallel_attn
+ and self.config.num_ln_in_parallel_attn == 1
+ ):
+ mlp_layernorm_out = attention_layernorm_out
+
+ # MLP.
+ mlp_output = self.mlp(mlp_layernorm_out)
+
+ if self.config.new_decoder_architecture or self.config.parallel_attn:
+ mlp_output += attention_output
+
+ output = dropout_add(mlp_output, residual, self.config.hidden_dropout, training=self.training)
+
+ return output, attn_weights
+
+
+@auto_docstring
+class FalconPreTrainedModel(PreTrainedModel):
+ config: FalconConfig
+ base_model_prefix = "transformer"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["FalconDecoderLayer"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _can_compile_fullgraph = True
+
+ @torch.no_grad()
+ def _init_weights(self, module: nn.Module):
+ """Initialize the weights."""
+ super()._init_weights(module)
+ if isinstance(module, FalconLinear):
+ init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+ if module.bias is not None:
+ init.zeros_(module.bias)
+
+ # Adapted from transformers.modeling_utils.PreTrainedModel._check_and_enable_sdpa
+ @classmethod
+ def _check_and_enable_sdpa(cls, config, hard_check_only: bool = False):
+ _is_bettertransformer = getattr(cls, "use_bettertransformer", False)
+ if _is_bettertransformer:
+ return config
+
+ if not hard_check_only:
+ config._attn_implementation = "sdpa"
+ return config
+
+
+@auto_docstring
+class FalconModel(FalconPreTrainedModel):
+ def __init__(self, config: FalconConfig):
+ super().__init__(config)
+
+ self.embed_dim = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.use_alibi = config.alibi
+
+ # Embedding + LN Embedding
+ self.word_embeddings = nn.Embedding(config.vocab_size, self.embed_dim)
+
+ # Transformer blocks
+ self.h = nn.ModuleList([FalconDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])
+
+ # Final Layer Norm
+ self.ln_f = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
+ self.gradient_checkpointing = False
+ self.rotary_emb = FalconRotaryEmbedding(config=config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.word_embeddings
+
+ def set_input_embeddings(self, new_embeddings: torch.Tensor):
+ self.word_embeddings = new_embeddings
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor, ...] | BaseModelOutputWithPastAndCrossAttentions:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+ `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
+ (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
+
+ If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+ `input_ids`.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ """
+ 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.return_dict
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or 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
+
+ if inputs_embeds is None:
+ inputs_embeds = self.word_embeddings(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ # Compute alibi tensor: check build_alibi_tensor documentation
+ alibi = None
+ past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+ batch_size, seq_length, _ = inputs_embeds.shape
+ if self.use_alibi:
+ mask = (
+ torch.ones(
+ (batch_size, seq_length + past_key_values_length), device=inputs_embeds.device, dtype=torch.long
+ )
+ if attention_mask is None
+ else attention_mask
+ )
+ alibi = build_alibi_tensor(mask, self.num_heads, dtype=inputs_embeds.dtype)
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ # Force mask creation for alibi
+ and_mask_function=lambda *args: torch.tensor(True, dtype=torch.bool),
+ )
+ if alibi is not None and causal_mask is not None and causal_mask.ndim == 4:
+ min_dtype = torch.finfo(inputs_embeds.dtype).min
+
+ # Only using non-bool mask for alibi
+ if causal_mask.dtype == torch.bool:
+ causal_mask = torch.where(
+ causal_mask, torch.tensor(0.0, device=causal_mask.device, dtype=inputs_embeds.dtype), min_dtype
+ )
+
+ # We take care to integrate alibi bias in the causal_mask here
+ alibi = alibi.reshape(batch_size, -1, *alibi.shape[1:])
+ causal_mask = torch.masked_fill(
+ alibi / math.sqrt(self.config.hidden_size // self.num_heads),
+ causal_mask < -1,
+ min_dtype,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ all_self_attentions = () if output_attentions else None
+ all_hidden_states = () if output_hidden_states else None
+
+ for i, block in enumerate(self.h):
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ outputs = block(
+ hidden_states,
+ layer_past=past_key_values,
+ attention_mask=causal_mask,
+ position_ids=position_ids,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ alibi=alibi,
+ position_embeddings=position_embeddings,
+ )
+
+ hidden_states = outputs[0]
+ if output_attentions:
+ all_self_attentions = all_self_attentions + (outputs[1],)
+
+ # Add last hidden state
+ hidden_states = self.ln_f(hidden_states)
+
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ if not return_dict:
+ return tuple(
+ v for v in [hidden_states, past_key_values, all_hidden_states, all_self_attentions] if v is not None
+ )
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The Falcon Model transformer with a language modeling head on top (linear layer with weights tied to the input embeddings).
+ """
+)
+class FalconForCausalLM(FalconPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "transformer.word_embeddings.weight"}
+
+ def __init__(self, config: FalconConfig):
+ super().__init__(config)
+ self.transformer = FalconModel(config)
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def set_output_embeddings(self, new_embeddings: torch.Tensor):
+ self.lm_head = new_embeddings
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs,
+ ) -> tuple[torch.Tensor] | CausalLMOutputWithCrossAttentions:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+ `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
+ (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
+
+ If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+ `input_ids`.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
+ `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
+ are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
+ """
+
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ transformer_outputs = self.transformer(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ hidden_states = transformer_outputs[0]
+
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ lm_logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ lm_logits,
+ labels,
+ vocab_size=self.config.vocab_size,
+ **kwargs,
+ )
+
+ if not return_dict:
+ output = (lm_logits,) + transformer_outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return CausalLMOutputWithCrossAttentions(
+ loss=loss,
+ logits=lm_logits,
+ past_key_values=transformer_outputs.past_key_values,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The Falcon Model transformer with a sequence classification head on top (linear layer).
+
+ [`FalconForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+ (e.g. GPT-1) do.
+
+ Since it does classification on the last token, it requires to know the position of the last token. If a
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+ each row of the batch).
+ """
+)
+class FalconForSequenceClassification(FalconPreTrainedModel):
+ def __init__(self, config: FalconConfig):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.transformer = FalconModel(config)
+ self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor] | SequenceClassifierOutputWithPast:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+ `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
+ (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
+
+ If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+ `input_ids`.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ transformer_outputs = self.transformer(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = transformer_outputs[0]
+ logits = self.score(hidden_states)
+
+ if input_ids is not None:
+ batch_size = input_ids.shape[0]
+ else:
+ batch_size = inputs_embeds.shape[0]
+
+ if self.config.pad_token_id is None and batch_size != 1:
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+ if self.config.pad_token_id is None:
+ last_non_pad_token = -1
+ elif input_ids is not None:
+ # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
+ non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+ token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+ last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+ else:
+ last_non_pad_token = -1
+ logger.warning_once(
+ f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+ "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+ )
+
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]
+
+ loss = None
+ if labels is not None:
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(pooled_logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(pooled_logits, labels)
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(pooled_logits, labels)
+ if not return_dict:
+ output = (pooled_logits,) + transformer_outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return SequenceClassifierOutputWithPast(
+ loss=loss,
+ logits=pooled_logits,
+ past_key_values=transformer_outputs.past_key_values,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+
+@auto_docstring
+class FalconForTokenClassification(FalconPreTrainedModel):
+ def __init__(self, config: FalconConfig):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.transformer = FalconModel(config)
+ if getattr(config, "classifier_dropout", None) is not None:
+ classifier_dropout = config.classifier_dropout
+ elif getattr(config, "hidden_dropout", None) is not None:
+ classifier_dropout = config.hidden_dropout
+ else:
+ classifier_dropout = 0.1
+ self.dropout = nn.Dropout(classifier_dropout)
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor] | TokenClassifierOutput:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+ `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
+ (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
+
+ If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+ `input_ids`.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ transformer_outputs = self.transformer(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = transformer_outputs[0]
+ hidden_states = self.dropout(hidden_states)
+ logits = self.classifier(hidden_states)
+
+ loss = None
+ if labels is not None:
+ batch_size, seq_length = labels.shape
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(
+ logits.view(batch_size * seq_length, self.num_labels), labels.view(batch_size * seq_length)
+ )
+
+ if not return_dict:
+ output = (logits,) + transformer_outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return TokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=transformer_outputs.hidden_states,
+ attentions=transformer_outputs.attentions,
+ )
+
+
+@auto_docstring
+class FalconForQuestionAnswering(FalconPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.transformer = FalconModel(config)
+ self.qa_outputs = nn.Linear(config.hidden_size, 2)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ start_positions: torch.LongTensor | None = None,
+ end_positions: torch.LongTensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | QuestionAnsweringModelOutput:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+ `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values.get_seq_length()`
+ (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
+
+ If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+ `input_ids`.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.transformer(
+ input_ids,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ sequence_output = outputs[0]
+
+ logits = self.qa_outputs(sequence_output)
+ start_logits, end_logits = logits.split(1, dim=-1)
+ start_logits = start_logits.squeeze(-1).contiguous()
+ end_logits = end_logits.squeeze(-1).contiguous()
+
+ total_loss = None
+ if start_positions is not None and end_positions is not None:
+ # If we are on multi-GPU, split add a dimension
+ if len(start_positions.size()) > 1:
+ start_positions = start_positions.squeeze(-1)
+ if len(end_positions.size()) > 1:
+ end_positions = end_positions.squeeze(-1)
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
+ ignored_index = start_logits.size(1)
+ start_positions = start_positions.clamp(0, ignored_index)
+ end_positions = end_positions.clamp(0, ignored_index)
+
+ loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+ start_loss = loss_fct(start_logits, start_positions)
+ end_loss = loss_fct(end_logits, end_positions)
+ total_loss = (start_loss + end_loss) / 2
+
+ if not return_dict:
+ output = (start_logits, end_logits) + outputs[2:]
+ return ((total_loss,) + output) if total_loss is not None else output
+
+ return QuestionAnsweringModelOutput(
+ loss=total_loss,
+ start_logits=start_logits,
+ end_logits=end_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "FalconForCausalLM",
+ "FalconModel",
+ "FalconPreTrainedModel",
+ "FalconForSequenceClassification",
+ "FalconForTokenClassification",
+ "FalconForQuestionAnswering",
+]
diff --git a/third_party/transformers/src/transformers/models/falcon_h1/__init__.py b/third_party/transformers/src/transformers/models/falcon_h1/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..9749c5e1e982cd1ac49db2d560238e4d45907c62
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/falcon_h1/__init__.py
@@ -0,0 +1,27 @@
+# Copyright 2025 TII and the HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_falcon_h1 import *
+ from .modeling_falcon_h1 import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/falcon_h1/configuration_falcon_h1.py b/third_party/transformers/src/transformers/models/falcon_h1/configuration_falcon_h1.py
new file mode 100644
index 0000000000000000000000000000000000000000..bb33b919be0a30b6ca9dfc4a597e230a47e7fccd
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/falcon_h1/configuration_falcon_h1.py
@@ -0,0 +1,139 @@
+# Copyright 2025 TII and the HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""FalconH1 model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="tiiuae/Falcon-H1-1.5B-Deep-Instruct")
+@strict
+class FalconH1Config(PreTrainedConfig):
+ r"""
+ num_logits_to_keep (`int` or `None`, *optional*, defaults to 1):
+ Number of prompt logits to calculate during generation. If `None`, all logits will be calculated. If an
+ integer value, only last `num_logits_to_keep` logits will be calculated. Default is 1 because only the
+ logits of the last prompt token are needed for generation. For long sequences, the logits for the entire
+ sequence may use a lot of memory so, setting `num_logits_to_keep=1` will reduce memory footprint
+ significantly.
+ projectors_bias (`bool`, *optional*, defaults to `False`):
+ Flag indicating whether or not to use bias in the input and output projections (["in_proj", "out_proj"]) of the attention block
+ lm_head_multiplier (`float`, *optional*, defaults to 1.0):
+ The multiplier for the LM head. This is used to scale the output of the LM head.
+ embedding_multiplier (`float`, *optional*, defaults to 1.0):
+ The multiplier for the embedding layer. This is used to scale the output of the embedding layer.
+ mlp_multipliers (`list[float]`, *optional*):
+ The multipliers for the MLP layers. This is used to scale the output of the MLP layers. The first value is
+ the multiplier of gate layer, the second value is the multiplier of the down_proj layer.
+ key_multiplier (`float`, *optional*):
+ The multiplier for the key layer. This is used to scale the output of the key layer.
+ attention_out_multiplier (`float`, *optional*):
+ The multiplier for the attention output layer. This is used to scale the output of the attention output
+ attention_in_multiplier (`float`, *optional*):
+ The multiplier for the attention input layer. This is used to scale the output of the attention input layer.
+ ssm_multipliers (`list[float]`, *optional*):
+ The multipliers for the SSM layers. This is used to scale the output of the SSM layers.
+ ssm_in_multiplier (`float`, *optional*):
+ The multiplier for the SSM input layer. This is used to scale the output of the SSM input layer.
+ ssm_out_multiplier (`float`, *optional*):
+ The multiplier for the SSM output layer. This is used to scale the output of the SSM output layer.
+ """
+
+ model_type = "falcon_h1"
+ attribute_map = {"layer_types": "layers_block_type"}
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ vocab_size: int = 128000
+ tie_word_embeddings: bool = False
+ hidden_size: int = 4096
+ intermediate_size: int = 14336
+ num_hidden_layers: int = 32
+ num_attention_heads: int = 32
+ num_key_value_heads: int | None = 8
+ hidden_act: str = "silu"
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-5
+ use_cache: bool | None = True
+ num_logits_to_keep: int | None = 1
+ pad_token_id: int | None = 0
+ bos_token_id: int | None = 1
+ eos_token_id: int | list[int] | None = 2
+ max_position_embeddings: int = 8192
+ attention_dropout: float | int | None = 0.0
+ mamba_d_ssm: int | None = 1024
+ mamba_n_heads: int | None = 128
+ mamba_d_head: str | int | None = "auto"
+ mamba_n_groups: int | None = 1
+ mamba_d_state: int | None = 256
+ mamba_d_conv: int | None = 4
+ mamba_expand: int | None = 2
+ mamba_chunk_size: int | None = 256
+ mamba_conv_bias: bool | None = True
+ mamba_proj_bias: bool | None = False
+ mamba_norm_before_gate: bool | None = True
+ mamba_rms_norm: bool | None = False
+ time_step_min: float | None = 0.001
+ time_step_max: float | None = 0.1
+ time_step_limit: list[float, float] | tuple[float, float] | None = (0.0, float("inf"))
+ projectors_bias: bool | None = False
+ rope_parameters: RopeParameters | dict | None = None
+ lm_head_multiplier: float | None = 1.0
+ embedding_multiplier: float | None = 1.0
+ mlp_multipliers: list[float] | None = None
+ key_multiplier: float | None = 1.0
+ attention_out_multiplier: float | None = 1.0
+ attention_in_multiplier: float | None = 1.0
+ ssm_multipliers: list[float] | None = None
+ ssm_in_multiplier: float | None = 1.0
+ ssm_out_multiplier: float | None = 1.0
+ attention_bias: bool = False
+ mlp_bias: bool = False
+
+ def __post_init__(self, **kwargs):
+ if self.num_key_value_heads is None:
+ self.num_key_value_heads = self.num_attention_heads
+
+ # for the mamba_v2, must satisfy the following
+ mamba_intermediate = self.mamba_expand * self.hidden_size if self.mamba_d_ssm is None else self.mamba_d_ssm
+ if self.mamba_d_head == "auto":
+ self.mamba_d_head = mamba_intermediate // self.mamba_n_heads
+
+ self.time_step_limit = tuple(self.time_step_limit) if self.time_step_limit is not None else None
+ if self.mlp_multipliers is None:
+ self.mlp_multipliers = [1.0, 1.0]
+
+ if self.ssm_multipliers is None:
+ self.ssm_multipliers = [1.0, 1.0, 1.0, 1.0, 1.0]
+
+ super().__post_init__(**kwargs)
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ mamba_intermediate = self.mamba_expand * self.hidden_size if self.mamba_d_ssm is None else self.mamba_d_ssm
+
+ if mamba_intermediate % self.mamba_n_heads != 0:
+ raise ValueError("mamba_n_heads must divide mamba_expand * hidden_size")
+
+ if self.mamba_d_head * self.mamba_n_heads != mamba_intermediate:
+ raise ValueError("The dimensions for the Mamba head state do not match the model intermediate_size")
+
+ @property
+ def layers_block_type(self):
+ return ["hybrid" for i in range(self.num_hidden_layers)]
+
+
+__all__ = ["FalconH1Config"]
diff --git a/third_party/transformers/src/transformers/models/falcon_h1/convert_mamba_ssm_checkpoint.py b/third_party/transformers/src/transformers/models/falcon_h1/convert_mamba_ssm_checkpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..ffcb58364f0b77c550aad33e94cd16f2230e7160
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/falcon_h1/convert_mamba_ssm_checkpoint.py
@@ -0,0 +1,148 @@
+# Copyright 2025 TII and the HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""This script can be used to convert checkpoints provided in the `mamba_ssm` library into the format provided in HuggingFace `transformers`. It depends on the `mamba2_ssm` package to be installed."""
+
+import argparse
+
+import torch
+
+from transformers import AutoModelForCausalLM, AutoTokenizer, FalconH1Config, FalconH1ForCausalLM
+
+
+CONVERSION_MAPPING = {
+ "backbone": "model",
+ "embeddings": "embed_tokens",
+ "mixer.": "",
+ "mixer_ssm": "mamba",
+ "mixer_attn": "self_attn",
+ "mlp.": "feed_forward.",
+ "mlp_norm": "pre_ff_layernorm",
+ "ssm_proj": "mamba.in_proj",
+ "attn_out_proj": "o_proj",
+ ".norm.": ".input_layernorm.",
+ ".mamba.input_layernorm.": ".mamba.norm.",
+ ".ssm_out_proj.": ".mamba.out_proj.",
+ "norm_f": "final_layernorm",
+}
+
+
+def convert_falcon_h1_to_hf(input_model_path, output_path):
+ tokenizer = AutoTokenizer.from_pretrained(input_model_path)
+
+ model = AutoModelForCausalLM.from_pretrained(input_model_path, dtype=torch.bfloat16, trust_remote_code=True)
+
+ intermediate_size = int(model.config.expansion_factor * model.config.hidden_size)
+
+ if intermediate_size % 2 != 0:
+ intermediate_size = intermediate_size + (intermediate_size % 2)
+
+ new_config = FalconH1Config(
+ vocab_size=model.config.vocab_size,
+ tie_word_embeddings=model.config.tie_word_embeddings,
+ hidden_size=model.config.hidden_size,
+ intermediate_size=intermediate_size,
+ mamba_d_state=model.config.state_size,
+ num_hidden_layers=model.config.num_hidden_layers,
+ mamba_use_mlp=model.config.use_mlp,
+ rms_norm_eps=model.config.layer_norm_epsilon,
+ pad_token_id=model.config.pad_token_id,
+ eos_token_id=model.config.eos_token_id,
+ mamba_expand=model.config.expand,
+ mamba_d_conv=model.config.conv_kernel,
+ mamba_n_groups=model.config.n_groups,
+ mamba_n_heads=model.config.num_heads,
+ mamba_norm_before_gate=model.config.norm_before_gate,
+ mamba_rms_norm=model.config.rms_norm,
+ mamba_d_ssm=model.config.d_ssm,
+ attention_bias=model.config.use_bias,
+ projectors_bias=model.config.use_bias,
+ mamba_conv_bias=model.config.use_conv_bias,
+ hidden_act=model.config.hidden_act,
+ use_cache=model.config.use_cache,
+ mamba_chunk_size=model.config.chunk_size,
+ num_attention_heads=model.config.num_heads_mha,
+ num_key_value_heads=model.config.num_key_value_heads,
+ head_dim=model.config.head_dim_mha,
+ lm_head_multiplier=model.config.lm_head_multiplier,
+ embedding_multiplier=model.config.embedding_multiplier,
+ mlp_multipliers=model.config.mlp_multipliers,
+ key_multiplier=model.config.key_multiplier,
+ attention_out_multiplier=model.config.attention_out_multiplier,
+ attention_in_multiplier=model.config.attention_in_multiplier,
+ ssm_multipliers=model.config.ssm_multipliers,
+ ssm_in_multiplier=model.config.ssm_in_multiplier,
+ ssm_out_multiplier=model.config.ssm_out_multiplier,
+ rope_theta=model.config.rope_theta,
+ )
+
+ old_state_dict = model.state_dict()
+ new_state_dict = {}
+
+ for old_key, old_value in old_state_dict.items():
+ new_key = old_key
+ for conversion_key, conversion_value in CONVERSION_MAPPING.items():
+ if conversion_key in old_key:
+ new_key = new_key.replace(conversion_key, conversion_value)
+
+ if "mamba.input_layernorm" in new_key:
+ new_key = new_key.replace("mamba.input_layernorm", "mamba.norm")
+
+ # Special processing for attention layers
+ if "self_attn.attn_proj" in new_key:
+ num_heads = new_config.num_attention_heads
+ num_kv_heads = new_config.num_key_value_heads
+ head_dim = new_config.head_dim
+ q_proj, k_proj, v_proj = old_value.split(
+ [
+ num_heads * head_dim,
+ num_kv_heads * head_dim,
+ num_kv_heads * head_dim,
+ ],
+ dim=0,
+ )
+ new_state_dict[new_key.replace("attn_proj", "q_proj")] = q_proj
+ new_state_dict[new_key.replace("attn_proj", "k_proj")] = k_proj
+ new_state_dict[new_key.replace("attn_proj", "v_proj")] = v_proj
+ else:
+ new_state_dict[new_key] = old_value
+
+ with torch.device("meta"):
+ new_model = FalconH1ForCausalLM(new_config)
+
+ del model
+
+ new_model.load_state_dict(new_state_dict, strict=True, assign=True)
+
+ new_model.save_pretrained(output_path)
+ tokenizer.save_pretrained(output_path)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "-i",
+ "--mamba_ssm_checkpoint_directory",
+ type=str,
+ required=True,
+ help="Path to a directory containing the `pytorch_model.bin` mamba_ssm checkpoint file to be converted.",
+ )
+ parser.add_argument(
+ "-o", "--output_dir", type=str, required=True, help="Path to directory to save the converted output model to."
+ )
+ args = parser.parse_args()
+
+ convert_falcon_h1_to_hf(
+ args.mamba_ssm_checkpoint_directory,
+ args.output_dir,
+ )
diff --git a/third_party/transformers/src/transformers/models/falcon_h1/modeling_falcon_h1.py b/third_party/transformers/src/transformers/models/falcon_h1/modeling_falcon_h1.py
new file mode 100644
index 0000000000000000000000000000000000000000..37b5da9df4b3e46fd93c6b8bde95286856adb005
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/falcon_h1/modeling_falcon_h1.py
@@ -0,0 +1,1265 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/falcon_h1/modular_falcon_h1.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_falcon_h1.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 Technology Innovation Institute 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.
+
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
+from ...integrations.hub_kernels import lazy_load_kernel
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
+from ...utils.generic import maybe_autocast, merge_with_config_defaults
+from ...utils.import_utils import resolve_internal_import
+from ...utils.output_capturing import capture_outputs
+from .configuration_falcon_h1 import FalconH1Config
+
+
+logger = logging.get_logger(__name__)
+
+
+class FalconH1RotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: FalconH1Config, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: FalconH1Config | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+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)
+
+
+@use_kernel_func_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+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)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class FalconH1Attention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: FalconH1Config, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = True
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+ )
+ self.key_multiplier = config.key_multiplier
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: torch.Tensor | None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) * self.key_multiplier
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class FalconH1RMSNormGated(torch.nn.Module):
+ def __init__(self, hidden_size, eps=1e-6, n_groups=1, norm_before_gate=True):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+ self.n_groups = n_groups
+ self.norm_before_gate = norm_before_gate
+
+ def forward(self, hidden_states, gate=None):
+ input_dtype = hidden_states.dtype
+
+ if not self.norm_before_gate and gate is not None:
+ hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+
+ if len(hidden_states.shape) == 3:
+ batch_size, seq_len, dim = hidden_states.shape
+ else:
+ batch_size, dim = hidden_states.shape
+ seq_len = 1
+ hidden_states = hidden_states.to(torch.float32)
+
+ hidden_states = hidden_states.view(batch_size, seq_len, self.n_groups, int(dim // self.n_groups))
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
+
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+
+ hidden_states = self.weight.view(self.n_groups, int(dim // self.n_groups)) * hidden_states
+ hidden_states = hidden_states.view(batch_size, seq_len, dim)
+
+ if seq_len == 1:
+ hidden_states = hidden_states.squeeze(1)
+
+ if self.norm_before_gate and gate is not None:
+ hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+ return hidden_states.to(input_dtype)
+
+
+# Helper methods for segment sum computation
+
+
+def pad_tensor_by_size(input_tensor: torch.Tensor, pad_size: int):
+ """
+ Padding x tensor with `pad_size` on the seq_len dim (dim=1)
+
+ Assumes that we only have tensors of either size 4 or 3
+ """
+ pad_shape = (0, 0, 0, 0, 0, pad_size, 0, 0) if len(input_tensor.shape) == 4 else (0, 0, 0, pad_size, 0, 0)
+
+ return torch.nn.functional.pad(input_tensor, pad_shape, mode="constant", value=0)
+
+
+def reshape_into_chunks(input_tensor, pad_size, chunk_size):
+ """
+ Padding input_tensor with `pad_size` on the seq_len dim (dim=1) and
+ simultaneously splitting it into chunk sequences.
+
+ Assumes that we only have tensors of either size 4 or 3
+ """
+ # [bsz, seq_len, ...] -> [bsz, seq_len multiple of chunk_size, ...]
+ input_tensor = pad_tensor_by_size(input_tensor, pad_size)
+
+ if len(input_tensor.shape) == 3:
+ # [bsz, seq_len multiple of chunk_size, num_heads] -> [bsz, -1, chunk_size, num_heads]
+ return input_tensor.reshape(input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2])
+ else:
+ # [bsz, seq_len multiple of chunk_size, num_heads, head_dim or state_size] -> [bsz, -1, chunk_size, num_heads, head_dim or state_size]
+ return input_tensor.reshape(
+ input_tensor.shape[0], -1, chunk_size, input_tensor.shape[2], input_tensor.shape[3]
+ )
+
+
+def segment_sum(input_tensor):
+ """
+ More stable segment sum calculation. Uses cumulative sums and masking instead of direct subtractions.
+ """
+ chunk_size = input_tensor.size(-1)
+ # 1. expand input tensor to have an additional dimension and repeat along that dimension
+ # [..., chunk_size] -> [..., chunk_size, chunk_size]
+ input_tensor = input_tensor[..., None].expand(*input_tensor.size(), chunk_size)
+ # 2. create a lower triangular mask with the diagonal set to 0 to 0 out elements above diag
+ mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=-1)
+ input_tensor = input_tensor.masked_fill(~mask, 0)
+ # 3. compute actual cumsum
+ tensor_segsum = torch.cumsum(input_tensor, dim=-2)
+
+ # 4. apply mask to keep only the lower triangular part of the cumulative sum result (incl diagonal this time)
+ mask = torch.tril(torch.ones(chunk_size, chunk_size, device=input_tensor.device, dtype=torch.bool), diagonal=0)
+ tensor_segsum = tensor_segsum.masked_fill(~mask, -torch.inf)
+ return tensor_segsum
+
+
+def apply_mask_to_padding_states(hidden_states, attention_mask):
+ """
+ Tunes out the hidden states for padding tokens, see https://github.com/state-spaces/mamba/issues/66
+ """
+ # NOTE: attention mask is a 2D boolean tensor
+ if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+ dtype = hidden_states.dtype
+ hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+
+ return hidden_states
+
+
+# Adapted from transformers.models.mamba2.modeling_mamba2.Mamba2Mixer
+class FalconH1Mixer(nn.Module):
+ """
+ FalconH1Mixer is identical to classic Mamba2 mixer classes but differs on two different things
+ - Users can pass custom intermediate_size through `config.mamba_d_ssm`
+ - The use of gated RMS normalization layer is optional
+ """
+
+ def __init__(self, config: FalconH1Config, layer_idx: int):
+ super().__init__()
+ self.num_heads = config.mamba_n_heads
+ self.hidden_size = config.hidden_size
+ self.ssm_state_size = config.mamba_d_state
+ self.conv_kernel_size = config.mamba_d_conv
+ self.intermediate_size = (
+ int(config.mamba_expand * self.hidden_size) if config.mamba_d_ssm is None else config.mamba_d_ssm
+ )
+ self.layer_idx = layer_idx
+ self.use_conv_bias = config.mamba_conv_bias
+ self.activation = config.hidden_act
+ self.act = ACT2FN[config.hidden_act]
+ self.use_bias = config.mamba_proj_bias
+
+ self.layer_norm_epsilon = config.rms_norm_eps
+ self.groups_time_state_size = config.mamba_n_groups * self.ssm_state_size
+
+ self.n_groups = config.mamba_n_groups
+ self.head_dim = config.mamba_d_head
+ self.chunk_size = config.mamba_chunk_size
+
+ self.time_step_limit = config.time_step_limit
+ self.time_step_min = config.time_step_min
+ self.time_step_max = config.time_step_max
+
+ self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size
+ self.conv1d = nn.Conv1d(
+ in_channels=self.conv_dim,
+ out_channels=self.conv_dim,
+ bias=config.mamba_conv_bias,
+ kernel_size=self.conv_kernel_size,
+ groups=self.conv_dim,
+ padding=self.conv_kernel_size - 1,
+ )
+
+ # projection of the input hidden states
+ projection_size = self.intermediate_size + self.conv_dim + self.num_heads
+ self.in_proj = nn.Linear(
+ self.hidden_size,
+ projection_size,
+ bias=self.use_bias,
+ )
+ # selective projection used to make dt, B and C input dependant
+
+ # time step projection (discretization)
+ # instantiate once and copy inv_dt in init_weights of PretrainedModel
+ self.dt_bias = nn.Parameter(torch.ones(self.num_heads))
+
+ # S4D real initialization. These are not discretized!
+ # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded
+ A = torch.arange(1, self.num_heads + 1)
+ self.A_log = nn.Parameter(torch.log(A))
+ self.mamba_rms_norm = config.mamba_rms_norm
+
+ if self.mamba_rms_norm:
+ self.norm = FalconH1RMSNormGated(
+ self.intermediate_size,
+ eps=self.layer_norm_epsilon,
+ n_groups=self.n_groups,
+ norm_before_gate=config.mamba_norm_before_gate,
+ )
+ self.D = nn.Parameter(torch.ones(self.num_heads))
+
+ self.out_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=config.projectors_bias)
+
+ global causal_conv1d_update, causal_conv1d_fn
+ causal_conv1d = lazy_load_kernel("causal-conv1d")
+ causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", None)
+ causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", None)
+
+ global selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
+ mamba_ssm = lazy_load_kernel("mamba-ssm")
+ selective_state_update = resolve_internal_import(
+ mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
+ )
+ mamba_chunk_scan_combined = resolve_internal_import(
+ mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
+ )
+ mamba_split_conv1d_scan_combined = resolve_internal_import(
+ mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
+ )
+
+ global is_fast_path_available
+ is_fast_path_available = all(
+ (
+ selective_state_update,
+ mamba_chunk_scan_combined,
+ mamba_split_conv1d_scan_combined,
+ causal_conv1d_fn,
+ causal_conv1d_update,
+ )
+ )
+
+ if not is_fast_path_available:
+ logger.warning_once(
+ "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
+ " is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
+ " https://github.com/Dao-AILab/causal-conv1d"
+ )
+ else:
+ logger.warning_once("The fast path for FalconH1 will be used when running the model on a GPU")
+
+ self.zxbcdt_multipliers = config.ssm_multipliers
+ self.ssm_in_multiplier = config.ssm_in_multiplier
+
+ def cuda_kernels_forward(
+ self,
+ hidden_states: torch.Tensor,
+ cache_params: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ ):
+ # 1. Gated MLP's linear projection
+ hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
+ # Add Multipliers
+ hidden_states = hidden_states * self.ssm_in_multiplier
+ projected_states = self.in_proj(hidden_states)
+ projected_states = projected_states * self.mup_vector # ADD Mup Multipliers
+ d_to_remove = 2 * self.intermediate_size + 2 * self.n_groups * self.ssm_state_size + self.num_heads
+
+ # Set up dimensions for reshapes later
+ batch_size, seq_len, _ = hidden_states.shape
+ groups_time_state_size = self.n_groups * self.ssm_state_size
+
+ use_precomputed_states = (
+ cache_params is not None and cache_params.has_previous_state(self.layer_idx) and seq_len == 1
+ )
+
+ # getting projected states from cache if it exists
+ if use_precomputed_states:
+ d_mlp = (projected_states.squeeze(1).shape[-1] - d_to_remove) // 2
+
+ z0, x0, gate, hidden_states_B_C, dt = projected_states.squeeze(1).split(
+ [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
+ )
+
+ # 2. Convolution sequence transformation
+ hidden_states_B_C = causal_conv1d_update(
+ hidden_states_B_C,
+ cache_params.layers[self.layer_idx].conv_states,
+ self.conv1d.weight.squeeze(1),
+ self.conv1d.bias,
+ self.activation,
+ )
+
+ hidden_states, B, C = torch.split(
+ hidden_states_B_C,
+ [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+ dim=-1,
+ )
+
+ # 3. SSM transformation
+ A = -torch.exp(self.A_log.float()) # (nheads,)
+ A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+ dt = dt[:, :, None].expand(-1, -1, self.head_dim)
+ dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
+ D = self.D[:, None, ...].expand(-1, self.head_dim)
+ B = B.view(batch_size, self.n_groups, B.shape[1] // self.n_groups)
+ C = C.view(batch_size, self.n_groups, C.shape[1] // self.n_groups)
+ hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)
+ hidden_states = selective_state_update(
+ cache_params.layers[self.layer_idx].recurrent_states,
+ hidden_states_reshaped,
+ dt,
+ A,
+ B,
+ C,
+ D,
+ z=gate.view(batch_size, self.num_heads, self.head_dim) if not self.mamba_rms_norm else None,
+ dt_bias=dt_bias,
+ dt_softplus=True,
+ )
+ hidden_states = hidden_states.view(batch_size, self.num_heads * self.head_dim)
+
+ if self.mamba_rms_norm:
+ hidden_states = self.norm(hidden_states, gate)
+
+ if d_mlp > 0:
+ hidden_states = torch.cat([F.silu(z0) * x0, hidden_states], dim=-1)
+
+ # 4. Final linear projection
+ out = self.out_proj(hidden_states[:, None, ...])
+ # Fused calculations or step by step if no initialized cache is found
+ else:
+ A = -torch.exp(self.A_log.float()) # (num_heads) or (intermediate_size, state_size)
+ dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
+
+ # 2-4. Fused kernel for conv1d, SSM, and the final projection
+ if self.training and cache_params is None:
+ out = mamba_split_conv1d_scan_combined(
+ projected_states,
+ self.conv1d.weight.squeeze(1),
+ self.conv1d.bias,
+ self.dt_bias,
+ A,
+ D=self.D,
+ chunk_size=self.chunk_size,
+ seq_idx=None, # was seq_idx
+ activation=self.activation,
+ rmsnorm_weight=self.norm.weight if self.mamba_rms_norm else None,
+ rmsnorm_eps=self.norm.variance_epsilon if self.mamba_rms_norm else None,
+ outproj_weight=self.out_proj.weight,
+ outproj_bias=self.out_proj.bias,
+ headdim=self.head_dim,
+ ngroups=self.n_groups,
+ norm_before_gate=False,
+ return_final_states=False,
+ **dt_limit_kwargs,
+ )
+
+ else:
+ d_mlp = (
+ projected_states.shape[-1]
+ - 2 * self.intermediate_size
+ - 2 * self.n_groups * self.ssm_state_size
+ - self.num_heads
+ ) // 2
+ if attention_mask is not None:
+ projected_states = projected_states * attention_mask[..., None]
+ _, gate, hidden_states_B_C, dt = projected_states.split(
+ [
+ 2 * d_mlp,
+ self.intermediate_size,
+ self.conv_dim,
+ self.num_heads,
+ ],
+ dim=-1,
+ )
+
+ if cache_params is not None:
+ conv_states = F.pad(
+ hidden_states_B_C.permute(0, 2, 1),
+ (self.conv_kernel_size - hidden_states_B_C.shape[-2], 0),
+ )
+ conv_states = cache_params.update_conv_state(conv_states, self.layer_idx)
+
+ time_step = nn.functional.softplus(dt + self.dt_bias)
+ # 1D Convolution
+ if causal_conv1d_fn is None or self.activation not in ["silu", "swish"]:
+ hidden_states_B_C = self.act(
+ self.conv1d(hidden_states_B_C.transpose(1, 2)).transpose(1, 2)[:, :seq_len]
+ ) # (B, L, self.d_inner + 2 * ngroups * d_state)
+ else:
+ hidden_states_B_C = causal_conv1d_fn(
+ x=hidden_states_B_C.transpose(1, 2),
+ weight=self.conv1d.weight.squeeze(1),
+ bias=self.conv1d.bias,
+ activation=self.activation,
+ ).transpose(1, 2)[:, :seq_len]
+
+ hidden_states, B, C = torch.split(
+ hidden_states_B_C,
+ [
+ self.intermediate_size,
+ groups_time_state_size,
+ groups_time_state_size,
+ ],
+ dim=-1,
+ )
+
+ if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+ # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66
+ dtype = hidden_states.dtype
+ hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+ # This is a hack to make sure multi-GPU inference works with HF accelerate
+ # see: https://github.com/Dao-AILab/flash-attention/issues/523 for more details
+ with torch.cuda.device(hidden_states.device):
+ scan_output, ssm_state = mamba_chunk_scan_combined(
+ hidden_states.view(batch_size, seq_len, -1, self.head_dim),
+ time_step,
+ A,
+ B.view(batch_size, seq_len, self.n_groups, -1),
+ C.view(batch_size, seq_len, self.n_groups, -1),
+ chunk_size=self.chunk_size,
+ D=self.D,
+ z=None,
+ seq_idx=None,
+ return_final_states=True,
+ **dt_limit_kwargs,
+ )
+ if ssm_state is not None and cache_params is not None:
+ ssm_state = cache_params.update_recurrent_state(ssm_state, self.layer_idx)
+ scan_output = scan_output.view(batch_size, seq_len, -1)
+ # Multiply "gate" branch and apply extra normalization layer
+ if self.mamba_rms_norm:
+ out = self.norm(scan_output, gate)
+ else:
+ out = scan_output * torch.nn.functional.silu(gate)
+ out = self.out_proj(out)
+ return out
+
+ # fmt: off
+ def torch_forward(
+ self,
+ input_states,
+ cache_params: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ ):
+ batch_size, seq_len, _ = input_states.shape
+ dtype = input_states.dtype
+
+ # 1. Gated MLP's linear projection
+ input_states = apply_mask_to_padding_states(input_states, attention_mask)
+ # Add Multipliers
+ input_states = input_states * self.ssm_in_multiplier
+ projected_states = self.in_proj(input_states)
+ projected_states = projected_states * self.mup_vector # ADD Mup Multipliers
+ gate, hidden_states_B_C, dt = projected_states.split([
+ self.intermediate_size, self.conv_dim, self.num_heads
+ ], dim=-1)
+ hidden_states_B_C = hidden_states_B_C.transpose(1,2)
+
+ use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) and seq_len == 1
+
+ # 2. Convolution sequence transformation
+ if use_precomputed_states:
+ conv_states = cache_params.update_conv_state(hidden_states_B_C, self.layer_idx)
+ # We need to guarantee that anything regarding the cache is on the same device
+ conv_states = conv_states.to(device=self.conv1d.weight.device)
+
+ hidden_states_B_C = torch.sum(
+ conv_states * self.conv1d.weight.squeeze(1), dim=-1
+ )
+ if self.use_conv_bias:
+ hidden_states_B_C = hidden_states_B_C + self.conv1d.bias
+ hidden_states_B_C = self.act(hidden_states_B_C)
+ else:
+ # Init cache
+ if cache_params is not None:
+ conv_states = nn.functional.pad(
+ hidden_states_B_C, (self.conv_kernel_size - hidden_states_B_C.shape[-1], 0)
+ )
+ conv_states = cache_params.update_conv_state(conv_states, self.layer_idx)
+
+ hidden_states_B_C = self.act(self.conv1d(hidden_states_B_C)[..., :seq_len].transpose(1, 2))
+
+ hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+ hidden_states, B, C = torch.split(
+ hidden_states_B_C,
+ [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size],
+ dim=-1
+ )
+
+ # 3. SSM transformation
+ A = -torch.exp(self.A_log.float()) # [num_heads]
+ if use_precomputed_states:
+ # We need to guarantee that anything regarding the cache is on the same device
+ cache_device = cache_params.layers[self.layer_idx].recurrent_states.device
+
+ # Note: there is no need to pad parameter matrices here, as there is just one new token
+ # for batched generation
+ dt = dt[:, 0, :][:, None, ...]
+ dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim)
+ # [num_heads] -> [num_heads, head_dim]
+ dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim)
+
+ dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))
+ dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1])
+ A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+ # [bsz, num_heads, head_dim, state_size]
+ dA = (torch.exp(dt[..., None] * A)).to(device=cache_device)
+
+ # Discretize B
+ # [bsz, n_groups * state_size] -> [bsz, n_groups, 1, state_size] ->
+ # -> [bsz, n_groups, group to head repetition factor, state_size] -> [bsz, num_heads, state_size]
+ B = B.reshape(batch_size, self.n_groups, -1)[..., None, :]
+ B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()
+ B = B.reshape(batch_size, -1, B.shape[-1])
+ # [bsz, num_heads, head_dim, state_size]
+ dB = dt[..., None] * B[..., None, :]
+
+ # Discretize x into dB
+ # [bsz, intermediate_size] -> [bsz, num_heads, head_dim]
+ hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
+ dBx = (dB * hidden_states[..., None]).to(device=cache_device)
+
+ # State calculation
+ ssm_states = cache_params.layers[self.layer_idx].recurrent_states * dA + dBx
+ ssm_states = cache_params.update_recurrent_state(ssm_states, self.layer_idx)
+
+ # Subsequent output
+ # [bsz, n_groups * state_size] -> [bsz, num_heads, state_size]
+ C = C.reshape(batch_size, self.n_groups, -1)[..., None, :]
+ C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous()
+ C = C.reshape(batch_size, -1, C.shape[-1])
+ # [bsz, num_heads, head_dim]
+
+ ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) # Shape: [b, h, d, n]
+ # Reshape ssm_states to merge the first two dimensions
+ ssm_states_reshaped = ssm_states.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size) # Shape: [b*h, d, n]
+ C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) # Shape: [b*h, n, 1]
+ y = torch.bmm(ssm_states_reshaped, C_reshaped)
+ y = y.view(batch_size, self.num_heads, self.head_dim)
+
+ # D skip connection
+ # [num_heads] -> [num_heads, head_dim]
+ D = self.D[..., None].expand(self.D.shape[0], self.head_dim)
+ y = (y + hidden_states * D).to(y.dtype)
+
+ # [bsz, num_heads, head_dim] -> [bsz, 1, intermediate_size]
+ y = y.reshape(batch_size, -1)[:, None, ...]
+ else:
+ # begin ssd naive implementation without einsums
+ dt = nn.functional.softplus(dt + self.dt_bias)
+ dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1])
+ hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float()
+ B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()
+ C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()
+ B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
+ C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
+ pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size
+
+ D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size)
+
+ # Discretize x and A
+ hidden_states = hidden_states * dt[..., None]
+ A = A.to(hidden_states.dtype) * dt
+
+ # Rearrange into blocks/chunks
+ hidden_states, A, B, C = [reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C)]
+
+ # [bsz, -1, chunk_size, num_heads] -> [bsz, num_heads, -1, chunk_size]
+ A = A.permute(0, 3, 1, 2)
+ A_cumsum = torch.cumsum(A, dim=-1)
+
+ # 1. Compute the output for each intra-chunk (diagonal blocks)
+ # This is the analog of a causal mask
+ L = torch.exp(segment_sum(A))
+
+ # Contraction of C and B to get G (attention-weights like)
+ G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :] # shape: (b, c, l, s, h, n)
+ G = G_intermediate.sum(dim=-1) # shape: (b, c, l, s, h)
+
+ # Compute M, equivalent to applying attention mask to weights
+ M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]
+ M = M_intermediate.sum(dim=-1)
+
+ # Compute Y_diag (apply to values)
+ Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3)
+
+ # 2. Compute the state for each intra-chunk
+ # (right term of low-rank factorization of off-diagonal blocks; B terms)
+ decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)
+ B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None]
+ states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2)
+
+ # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
+ # (middle term of factorization of off-diag blocks; A terms)
+ previous_states = torch.zeros_like(states[:, :1])
+ states = torch.cat([previous_states, states], dim=1)
+ decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0))))
+ decay_chunk = decay_chunk.transpose(1, 3)
+ new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
+ states, ssm_state = new_states[:, :-1], new_states[:, -1]
+
+ # 4. Compute state -> output conversion per chunk
+ # (left term of low-rank factorization of off-diagonal blocks; C terms)
+ state_decay_out = torch.exp(A_cumsum)
+ C_times_states = (C[..., None, :] * states[:, :, None, ...])
+ state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1)
+ Y_off = (C_times_states.sum(-1) * state_decay_out_permuted[..., None])
+
+ # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
+ y = Y_diag + Y_off
+ # [bsz, -1, self.chunk_size, num_heads, head_dim] -> [bsz, (padded) seq_len, num_heads, head_dim]
+ y = y.reshape(batch_size, -1, self.num_heads, self.head_dim)
+
+ y = y + D_residual
+ # Cutting off padded chunks
+ if pad_size > 0:
+ y = y[:, :seq_len, :, :]
+ y = y.reshape(batch_size, seq_len, -1)
+
+ # Init cache
+ if ssm_state is not None and cache_params is not None:
+ ssm_state = cache_params.update_recurrent_state(ssm_state, self.layer_idx)
+
+ if self.mamba_rms_norm:
+ scan_output = self.norm(y, gate)
+ else:
+ scan_output = y * torch.nn.functional.silu(gate)
+
+ # end ssd naive
+
+ # 4. Final linear projection
+ contextualized_states = self.out_proj(scan_output.to(dtype)) # [batch, seq_len, hidden_size]
+ return contextualized_states
+ # fmt: on
+
+ def forward(
+ self,
+ hidden_states,
+ cache_params: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs,
+ ):
+ if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling():
+ return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask)
+ dtype = hidden_states.dtype
+ if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+ # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66
+ hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+
+ return self.torch_forward(hidden_states, cache_params, attention_mask)
+
+
+class FalconH1MLP(nn.Module):
+ def __init__(self, config: FalconH1Config):
+ 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=config.mlp_bias)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)
+ self.act_fn = ACT2FN[config.hidden_act]
+ self.gate_multiplier, self.down_multiplier = config.mlp_multipliers
+
+ def forward(self, x):
+ y = self.up_proj(x) * self.act_fn(self.gate_proj(x) * self.gate_multiplier)
+ y = self.down_proj(y) * self.down_multiplier
+ return y
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class FalconH1RMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ FalconH1RMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class FalconH1DecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: FalconH1Config, layer_idx: int):
+ super().__init__()
+ self.feed_forward = FalconH1MLP(config)
+
+ head_dim = config.hidden_size // config.num_attention_heads
+ self.channels_attn = config.num_attention_heads * head_dim + 2 * config.num_key_value_heads * head_dim
+
+ self.mamba = FalconH1Mixer(config=config, layer_idx=layer_idx)
+
+ self.self_attn = FalconH1Attention(config, layer_idx)
+
+ self.attention_in_multiplier = config.attention_in_multiplier
+ self.ssm_out_multiplier = config.ssm_out_multiplier
+ self.attn_out_multiplier = config.attention_out_multiplier
+
+ self.input_layernorm = FalconH1RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.pre_ff_layernorm = FalconH1RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ mamba_attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs,
+ ) -> tuple[torch.FloatTensor]:
+ """
+ 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.
+ past_key_values (`Cache`, *optional*): cached past key and value projection states
+ 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`).
+ position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+ Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+ with `head_dim` being the embedding dimension of each attention head.
+ kwargs (`dict`, *optional*):
+ Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+ into the model
+ """
+
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ mamba_hidden_states = self.mamba(
+ hidden_states=hidden_states,
+ cache_params=past_key_values,
+ attention_mask=mamba_attention_mask,
+ )
+ mamba_hidden_states = mamba_hidden_states * self.ssm_out_multiplier
+
+ attention_hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states * self.attention_in_multiplier,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ attention_hidden_states = attention_hidden_states * self.attn_out_multiplier
+
+ hidden_states = mamba_hidden_states + attention_hidden_states
+
+ # residual connection after attention
+ hidden_states = residual + hidden_states
+
+ # feed-forward
+ residual = hidden_states
+ hidden_states = self.pre_ff_layernorm(hidden_states)
+ hidden_states = self.feed_forward(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return (hidden_states,)
+
+
+def compute_mup_vector(config):
+ """
+ Computes the MuP vector based on model configuration.
+
+ FalconH1 applies different MuP multiplier for each dimension of the hidden states.
+ The MuP vector is partitioned into chunks, and each chunk is multiplied with its
+ corresponding projected dimension.
+
+ Args:
+ config: FalconH1Config object
+
+ Returns:
+ torch.Tensor: The computed MuP vector
+ """
+ # We'll need some values from the config to compute the vector dimensions
+ intermediate_size = (
+ config.mamba_d_ssm if config.mamba_d_ssm is not None else int(config.mamba_expand * config.hidden_size)
+ )
+ groups_time_state_size = config.mamba_n_groups * config.mamba_d_state
+ num_heads = config.mamba_n_heads
+ zxbcdt_multipliers = config.ssm_multipliers
+
+ vector_shape = 2 * intermediate_size + 2 * groups_time_state_size + num_heads
+ mup_vector = torch.ones(1, 1, vector_shape)
+
+ # Apply multipliers to different sections of the vector
+ mup_vector[:, :, :intermediate_size] *= zxbcdt_multipliers[0]
+ mup_vector[:, :, intermediate_size : 2 * intermediate_size] *= zxbcdt_multipliers[1]
+ mup_vector[:, :, 2 * intermediate_size : 2 * intermediate_size + groups_time_state_size] *= zxbcdt_multipliers[2]
+ mup_vector[
+ :, :, 2 * intermediate_size + groups_time_state_size : 2 * intermediate_size + 2 * groups_time_state_size
+ ] *= zxbcdt_multipliers[3]
+ mup_vector[:, :, 2 * intermediate_size + 2 * groups_time_state_size :] *= zxbcdt_multipliers[4]
+
+ return mup_vector
+
+
+@auto_docstring
+class FalconH1PreTrainedModel(PreTrainedModel):
+ config: FalconH1Config
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["FalconH1DecoderLayer"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _is_stateful = True
+
+ _can_record_outputs = {
+ "hidden_states": FalconH1DecoderLayer,
+ "attentions": FalconH1Attention,
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, FalconH1Mixer):
+ init.ones_(module.dt_bias)
+ init.copy_(module.A_log, torch.log(torch.arange(1, module.num_heads + 1)))
+ init.ones_(module.D)
+ elif isinstance(module, FalconH1Model):
+ mup_vector = compute_mup_vector(module.config)
+ for layer in module.layers:
+ init.copy_(layer.mamba.mup_vector, mup_vector)
+
+
+@auto_docstring
+# Adapted from transformers.models.jamba.modeling_jamba.JambaModel
+class FalconH1Model(FalconH1PreTrainedModel):
+ def __init__(self, config: FalconH1Config):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ decoder_layers = []
+ for i in range(config.num_hidden_layers):
+ decoder_layers.append(FalconH1DecoderLayer(config, layer_idx=i))
+ self.layers = nn.ModuleList(decoder_layers)
+
+ self._attn_implementation = config._attn_implementation
+ self.final_layernorm = FalconH1RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = FalconH1RotaryEmbedding(config=config)
+
+ self.embedding_multiplier = config.embedding_multiplier
+ self.lm_head_multiplier = config.lm_head_multiplier
+
+ self.gradient_checkpointing = False
+ # Compute the MuP vector once and register it for all layers
+ mup_vector = compute_mup_vector(config)
+ for layer in self.layers:
+ layer.mamba.register_buffer("mup_vector", mup_vector.clone(), persistent=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids) * self.embedding_multiplier
+ hidden_states = inputs_embeds
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+ mamba_mask = self._update_mamba_mask(attention_mask, past_key_values)
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for decoder_layer in self.layers:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ mamba_attention_mask=mamba_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ hidden_states = self.final_layernorm(hidden_states)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+ def _update_mamba_mask(self, attention_mask, past_key_values):
+ """
+ No need for zeroing states when
+ 1. Cached forward
+ 2. Attending to all inputs
+ """
+ mamba_mask = attention_mask
+ if (past_key_values is not None and past_key_values.has_previous_state()) or (
+ attention_mask is not None and torch.all(attention_mask == 1)
+ ):
+ mamba_mask = None
+ return mamba_mask
+
+
+@auto_docstring
+class FalconH1ForCausalLM(FalconH1PreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+ _tp_plan = {"lm_head": "colwise_gather_output"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = FalconH1Model(config)
+ self.vocab_size = config.vocab_size
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs,
+ ) -> tuple | CausalLMOutputWithPast:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, FalconH1ForCausalLM
+
+ >>> model = FalconH1ForCausalLM.from_pretrained("...")
+ >>> tokenizer = AutoTokenizer.from_pretrained("...")
+
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+ ```"""
+ 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,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :]) * self.model.lm_head_multiplier
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ position_ids=None,
+ use_cache=True,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ kwargs["logits_to_keep"] = self.config.num_logits_to_keep
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ position_ids=position_ids,
+ use_cache=use_cache,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ return model_inputs
+
+
+__all__ = ["FalconH1Model", "FalconH1ForCausalLM", "FalconH1PreTrainedModel"]
diff --git a/third_party/transformers/src/transformers/models/falcon_h1/modular_falcon_h1.py b/third_party/transformers/src/transformers/models/falcon_h1/modular_falcon_h1.py
new file mode 100644
index 0000000000000000000000000000000000000000..75cbed28a64618dca067c225981eeb9fd4af5c2b
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/falcon_h1/modular_falcon_h1.py
@@ -0,0 +1,1014 @@
+# Copyright 2025 Technology Innovation Institute 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 FalconH1 model."""
+
+from collections.abc import Callable
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...integrations.hub_kernels import lazy_load_kernel
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, is_torchdynamo_compiling, logging
+from ...utils.generic import merge_with_config_defaults
+from ...utils.import_utils import resolve_internal_import
+from ...utils.output_capturing import capture_outputs
+from ..llama.modeling_llama import (
+ LlamaAttention,
+ LlamaForCausalLM,
+ LlamaMLP,
+ LlamaRMSNorm,
+ LlamaRotaryEmbedding,
+ apply_rotary_pos_emb,
+ eager_attention_forward,
+)
+from ..mamba2.modeling_mamba2 import (
+ MambaRMSNormGated,
+ apply_mask_to_padding_states,
+ pad_tensor_by_size,
+ reshape_into_chunks,
+ segment_sum,
+)
+from .configuration_falcon_h1 import FalconH1Config
+
+
+logger = logging.get_logger(__name__)
+
+
+class FalconH1RotaryEmbedding(LlamaRotaryEmbedding):
+ pass
+
+
+class FalconH1Attention(LlamaAttention):
+ def __init__(self, config: FalconH1Config, layer_idx: int):
+ super().__init__(config, layer_idx)
+ self.key_multiplier = config.key_multiplier
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: torch.Tensor | None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) * self.key_multiplier
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class FalconH1RMSNormGated(MambaRMSNormGated):
+ def __init__(self, hidden_size, eps=1e-6, n_groups=1, norm_before_gate=True):
+ super().__init__(hidden_size=hidden_size, eps=eps)
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+ self.n_groups = n_groups
+ self.norm_before_gate = norm_before_gate
+
+ def forward(self, hidden_states, gate=None):
+ input_dtype = hidden_states.dtype
+
+ if not self.norm_before_gate and gate is not None:
+ hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+
+ if len(hidden_states.shape) == 3:
+ batch_size, seq_len, dim = hidden_states.shape
+ else:
+ batch_size, dim = hidden_states.shape
+ seq_len = 1
+ hidden_states = hidden_states.to(torch.float32)
+
+ hidden_states = hidden_states.view(batch_size, seq_len, self.n_groups, int(dim // self.n_groups))
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
+
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+
+ hidden_states = self.weight.view(self.n_groups, int(dim // self.n_groups)) * hidden_states
+ hidden_states = hidden_states.view(batch_size, seq_len, dim)
+
+ if seq_len == 1:
+ hidden_states = hidden_states.squeeze(1)
+
+ if self.norm_before_gate and gate is not None:
+ hidden_states = hidden_states * F.silu(gate.to(torch.float32))
+ return hidden_states.to(input_dtype)
+
+
+# Adapted from transformers.models.mamba2.modeling_mamba2.Mamba2Mixer
+class FalconH1Mixer(nn.Module):
+ """
+ FalconH1Mixer is identical to classic Mamba2 mixer classes but differs on two different things
+ - Users can pass custom intermediate_size through `config.mamba_d_ssm`
+ - The use of gated RMS normalization layer is optional
+ """
+
+ def __init__(self, config: FalconH1Config, layer_idx: int):
+ super().__init__()
+ self.num_heads = config.mamba_n_heads
+ self.hidden_size = config.hidden_size
+ self.ssm_state_size = config.mamba_d_state
+ self.conv_kernel_size = config.mamba_d_conv
+ self.intermediate_size = (
+ int(config.mamba_expand * self.hidden_size) if config.mamba_d_ssm is None else config.mamba_d_ssm
+ )
+ self.layer_idx = layer_idx
+ self.use_conv_bias = config.mamba_conv_bias
+ self.activation = config.hidden_act
+ self.act = ACT2FN[config.hidden_act]
+ self.use_bias = config.mamba_proj_bias
+
+ self.layer_norm_epsilon = config.rms_norm_eps
+ self.groups_time_state_size = config.mamba_n_groups * self.ssm_state_size
+
+ self.n_groups = config.mamba_n_groups
+ self.head_dim = config.mamba_d_head
+ self.chunk_size = config.mamba_chunk_size
+
+ self.time_step_limit = config.time_step_limit
+ self.time_step_min = config.time_step_min
+ self.time_step_max = config.time_step_max
+
+ self.conv_dim = self.intermediate_size + 2 * self.n_groups * self.ssm_state_size
+ self.conv1d = nn.Conv1d(
+ in_channels=self.conv_dim,
+ out_channels=self.conv_dim,
+ bias=config.mamba_conv_bias,
+ kernel_size=self.conv_kernel_size,
+ groups=self.conv_dim,
+ padding=self.conv_kernel_size - 1,
+ )
+
+ # projection of the input hidden states
+ projection_size = self.intermediate_size + self.conv_dim + self.num_heads
+ self.in_proj = nn.Linear(
+ self.hidden_size,
+ projection_size,
+ bias=self.use_bias,
+ )
+ # selective projection used to make dt, B and C input dependant
+
+ # time step projection (discretization)
+ # instantiate once and copy inv_dt in init_weights of PretrainedModel
+ self.dt_bias = nn.Parameter(torch.ones(self.num_heads))
+
+ # S4D real initialization. These are not discretized!
+ # The core is to load them, compute the discrete states, then write the updated state. Keeps the memory bounded
+ A = torch.arange(1, self.num_heads + 1)
+ self.A_log = nn.Parameter(torch.log(A))
+ self.mamba_rms_norm = config.mamba_rms_norm
+
+ if self.mamba_rms_norm:
+ self.norm = FalconH1RMSNormGated(
+ self.intermediate_size,
+ eps=self.layer_norm_epsilon,
+ n_groups=self.n_groups,
+ norm_before_gate=config.mamba_norm_before_gate,
+ )
+ self.D = nn.Parameter(torch.ones(self.num_heads))
+
+ self.out_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=config.projectors_bias)
+
+ global causal_conv1d_update, causal_conv1d_fn
+ causal_conv1d = lazy_load_kernel("causal-conv1d")
+ causal_conv1d_update = getattr(causal_conv1d, "causal_conv1d_update", None)
+ causal_conv1d_fn = getattr(causal_conv1d, "causal_conv1d_fn", None)
+
+ global selective_state_update, mamba_chunk_scan_combined, mamba_split_conv1d_scan_combined
+ mamba_ssm = lazy_load_kernel("mamba-ssm")
+ selective_state_update = resolve_internal_import(
+ mamba_ssm, chained_path="ops.triton.selective_state_update.selective_state_update"
+ )
+ mamba_chunk_scan_combined = resolve_internal_import(
+ mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_chunk_scan_combined"
+ )
+ mamba_split_conv1d_scan_combined = resolve_internal_import(
+ mamba_ssm, chained_path="ops.triton.ssd_combined.mamba_split_conv1d_scan_combined"
+ )
+
+ global is_fast_path_available
+ is_fast_path_available = all(
+ (
+ selective_state_update,
+ mamba_chunk_scan_combined,
+ mamba_split_conv1d_scan_combined,
+ causal_conv1d_fn,
+ causal_conv1d_update,
+ )
+ )
+
+ if not is_fast_path_available:
+ logger.warning_once(
+ "The fast path is not available because one of `(selective_state_update, causal_conv1d_fn, causal_conv1d_update)`"
+ " is None. Falling back to the naive implementation. To install follow https://github.com/state-spaces/mamba/#installation and"
+ " https://github.com/Dao-AILab/causal-conv1d"
+ )
+ else:
+ logger.warning_once("The fast path for FalconH1 will be used when running the model on a GPU")
+
+ self.zxbcdt_multipliers = config.ssm_multipliers
+ self.ssm_in_multiplier = config.ssm_in_multiplier
+
+ def cuda_kernels_forward(
+ self,
+ hidden_states: torch.Tensor,
+ cache_params: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ ):
+ # 1. Gated MLP's linear projection
+ hidden_states = apply_mask_to_padding_states(hidden_states, attention_mask)
+ # Add Multipliers
+ hidden_states = hidden_states * self.ssm_in_multiplier
+ projected_states = self.in_proj(hidden_states)
+ projected_states = projected_states * self.mup_vector # ADD Mup Multipliers
+ d_to_remove = 2 * self.intermediate_size + 2 * self.n_groups * self.ssm_state_size + self.num_heads
+
+ # Set up dimensions for reshapes later
+ batch_size, seq_len, _ = hidden_states.shape
+ groups_time_state_size = self.n_groups * self.ssm_state_size
+
+ use_precomputed_states = (
+ cache_params is not None and cache_params.has_previous_state(self.layer_idx) and seq_len == 1
+ )
+
+ # getting projected states from cache if it exists
+ if use_precomputed_states:
+ d_mlp = (projected_states.squeeze(1).shape[-1] - d_to_remove) // 2
+
+ z0, x0, gate, hidden_states_B_C, dt = projected_states.squeeze(1).split(
+ [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1
+ )
+
+ # 2. Convolution sequence transformation
+ hidden_states_B_C = causal_conv1d_update(
+ hidden_states_B_C,
+ cache_params.layers[self.layer_idx].conv_states,
+ self.conv1d.weight.squeeze(1),
+ self.conv1d.bias,
+ self.activation,
+ )
+
+ hidden_states, B, C = torch.split(
+ hidden_states_B_C,
+ [self.intermediate_size, groups_time_state_size, groups_time_state_size],
+ dim=-1,
+ )
+
+ # 3. SSM transformation
+ A = -torch.exp(self.A_log.float()) # (nheads,)
+ A = A[:, None, ...][:, :, None].expand(-1, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+ dt = dt[:, :, None].expand(-1, -1, self.head_dim)
+ dt_bias = self.dt_bias[:, None, ...].expand(-1, self.head_dim)
+ D = self.D[:, None, ...].expand(-1, self.head_dim)
+ B = B.view(batch_size, self.n_groups, B.shape[1] // self.n_groups)
+ C = C.view(batch_size, self.n_groups, C.shape[1] // self.n_groups)
+ hidden_states_reshaped = hidden_states.view(batch_size, self.num_heads, self.head_dim)
+ hidden_states = selective_state_update(
+ cache_params.layers[self.layer_idx].recurrent_states,
+ hidden_states_reshaped,
+ dt,
+ A,
+ B,
+ C,
+ D,
+ z=gate.view(batch_size, self.num_heads, self.head_dim) if not self.mamba_rms_norm else None,
+ dt_bias=dt_bias,
+ dt_softplus=True,
+ )
+ hidden_states = hidden_states.view(batch_size, self.num_heads * self.head_dim)
+
+ if self.mamba_rms_norm:
+ hidden_states = self.norm(hidden_states, gate)
+
+ if d_mlp > 0:
+ hidden_states = torch.cat([F.silu(z0) * x0, hidden_states], dim=-1)
+
+ # 4. Final linear projection
+ out = self.out_proj(hidden_states[:, None, ...])
+ # Fused calculations or step by step if no initialized cache is found
+ else:
+ A = -torch.exp(self.A_log.float()) # (num_heads) or (intermediate_size, state_size)
+ dt_limit_kwargs = {} if self.time_step_limit == (0.0, float("inf")) else {"dt_limit": self.time_step_limit}
+
+ # 2-4. Fused kernel for conv1d, SSM, and the final projection
+ if self.training and cache_params is None:
+ out = mamba_split_conv1d_scan_combined(
+ projected_states,
+ self.conv1d.weight.squeeze(1),
+ self.conv1d.bias,
+ self.dt_bias,
+ A,
+ D=self.D,
+ chunk_size=self.chunk_size,
+ seq_idx=None, # was seq_idx
+ activation=self.activation,
+ rmsnorm_weight=self.norm.weight if self.mamba_rms_norm else None,
+ rmsnorm_eps=self.norm.variance_epsilon if self.mamba_rms_norm else None,
+ outproj_weight=self.out_proj.weight,
+ outproj_bias=self.out_proj.bias,
+ headdim=self.head_dim,
+ ngroups=self.n_groups,
+ norm_before_gate=False,
+ return_final_states=False,
+ **dt_limit_kwargs,
+ )
+
+ else:
+ d_mlp = (
+ projected_states.shape[-1]
+ - 2 * self.intermediate_size
+ - 2 * self.n_groups * self.ssm_state_size
+ - self.num_heads
+ ) // 2
+ if attention_mask is not None:
+ projected_states = projected_states * attention_mask[..., None]
+ _, gate, hidden_states_B_C, dt = projected_states.split(
+ [
+ 2 * d_mlp,
+ self.intermediate_size,
+ self.conv_dim,
+ self.num_heads,
+ ],
+ dim=-1,
+ )
+
+ if cache_params is not None:
+ conv_states = F.pad(
+ hidden_states_B_C.permute(0, 2, 1),
+ (self.conv_kernel_size - hidden_states_B_C.shape[-2], 0),
+ )
+ conv_states = cache_params.update_conv_state(conv_states, self.layer_idx)
+
+ time_step = nn.functional.softplus(dt + self.dt_bias)
+ # 1D Convolution
+ if causal_conv1d_fn is None or self.activation not in ["silu", "swish"]:
+ hidden_states_B_C = self.act(
+ self.conv1d(hidden_states_B_C.transpose(1, 2)).transpose(1, 2)[:, :seq_len]
+ ) # (B, L, self.d_inner + 2 * ngroups * d_state)
+ else:
+ hidden_states_B_C = causal_conv1d_fn(
+ x=hidden_states_B_C.transpose(1, 2),
+ weight=self.conv1d.weight.squeeze(1),
+ bias=self.conv1d.bias,
+ activation=self.activation,
+ ).transpose(1, 2)[:, :seq_len]
+
+ hidden_states, B, C = torch.split(
+ hidden_states_B_C,
+ [
+ self.intermediate_size,
+ groups_time_state_size,
+ groups_time_state_size,
+ ],
+ dim=-1,
+ )
+
+ if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+ # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66
+ dtype = hidden_states.dtype
+ hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+ # This is a hack to make sure multi-GPU inference works with HF accelerate
+ # see: https://github.com/Dao-AILab/flash-attention/issues/523 for more details
+ with torch.cuda.device(hidden_states.device):
+ scan_output, ssm_state = mamba_chunk_scan_combined(
+ hidden_states.view(batch_size, seq_len, -1, self.head_dim),
+ time_step,
+ A,
+ B.view(batch_size, seq_len, self.n_groups, -1),
+ C.view(batch_size, seq_len, self.n_groups, -1),
+ chunk_size=self.chunk_size,
+ D=self.D,
+ z=None,
+ seq_idx=None,
+ return_final_states=True,
+ **dt_limit_kwargs,
+ )
+ if ssm_state is not None and cache_params is not None:
+ ssm_state = cache_params.update_recurrent_state(ssm_state, self.layer_idx)
+ scan_output = scan_output.view(batch_size, seq_len, -1)
+ # Multiply "gate" branch and apply extra normalization layer
+ if self.mamba_rms_norm:
+ out = self.norm(scan_output, gate)
+ else:
+ out = scan_output * torch.nn.functional.silu(gate)
+ out = self.out_proj(out)
+ return out
+
+ # fmt: off
+ def torch_forward(
+ self,
+ input_states,
+ cache_params: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ ):
+ batch_size, seq_len, _ = input_states.shape
+ dtype = input_states.dtype
+
+ # 1. Gated MLP's linear projection
+ input_states = apply_mask_to_padding_states(input_states, attention_mask)
+ # Add Multipliers
+ input_states = input_states * self.ssm_in_multiplier
+ projected_states = self.in_proj(input_states)
+ projected_states = projected_states * self.mup_vector # ADD Mup Multipliers
+ gate, hidden_states_B_C, dt = projected_states.split([
+ self.intermediate_size, self.conv_dim, self.num_heads
+ ], dim=-1)
+ hidden_states_B_C = hidden_states_B_C.transpose(1,2)
+
+ use_precomputed_states = cache_params is not None and cache_params.has_previous_state(self.layer_idx) and seq_len == 1
+
+ # 2. Convolution sequence transformation
+ if use_precomputed_states:
+ conv_states = cache_params.update_conv_state(hidden_states_B_C, self.layer_idx)
+ # We need to guarantee that anything regarding the cache is on the same device
+ conv_states = conv_states.to(device=self.conv1d.weight.device)
+
+ hidden_states_B_C = torch.sum(
+ conv_states * self.conv1d.weight.squeeze(1), dim=-1
+ )
+ if self.use_conv_bias:
+ hidden_states_B_C = hidden_states_B_C + self.conv1d.bias
+ hidden_states_B_C = self.act(hidden_states_B_C)
+ else:
+ # Init cache
+ if cache_params is not None:
+ conv_states = nn.functional.pad(
+ hidden_states_B_C, (self.conv_kernel_size - hidden_states_B_C.shape[-1], 0)
+ )
+ conv_states = cache_params.update_conv_state(conv_states, self.layer_idx)
+
+ hidden_states_B_C = self.act(self.conv1d(hidden_states_B_C)[..., :seq_len].transpose(1, 2))
+
+ hidden_states_B_C = apply_mask_to_padding_states(hidden_states_B_C, attention_mask)
+ hidden_states, B, C = torch.split(
+ hidden_states_B_C,
+ [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size],
+ dim=-1
+ )
+
+ # 3. SSM transformation
+ A = -torch.exp(self.A_log.float()) # [num_heads]
+ if use_precomputed_states:
+ # We need to guarantee that anything regarding the cache is on the same device
+ cache_device = cache_params.layers[self.layer_idx].recurrent_states.device
+
+ # Note: there is no need to pad parameter matrices here, as there is just one new token
+ # for batched generation
+ dt = dt[:, 0, :][:, None, ...]
+ dt = dt.transpose(1, 2).expand(batch_size, dt.shape[-1], self.head_dim)
+ # [num_heads] -> [num_heads, head_dim]
+ dt_bias = self.dt_bias[..., None].expand(self.dt_bias.shape[0], self.head_dim)
+
+ dt = torch.nn.functional.softplus(dt + dt_bias.to(dt.dtype))
+ dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1])
+ A = A[..., None, None].expand(self.num_heads, self.head_dim, self.ssm_state_size).to(dtype=torch.float32)
+ # [bsz, num_heads, head_dim, state_size]
+ dA = (torch.exp(dt[..., None] * A)).to(device=cache_device)
+
+ # Discretize B
+ # [bsz, n_groups * state_size] -> [bsz, n_groups, 1, state_size] ->
+ # -> [bsz, n_groups, group to head repetition factor, state_size] -> [bsz, num_heads, state_size]
+ B = B.reshape(batch_size, self.n_groups, -1)[..., None, :]
+ B = B.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, B.shape[-1]).contiguous()
+ B = B.reshape(batch_size, -1, B.shape[-1])
+ # [bsz, num_heads, head_dim, state_size]
+ dB = dt[..., None] * B[..., None, :]
+
+ # Discretize x into dB
+ # [bsz, intermediate_size] -> [bsz, num_heads, head_dim]
+ hidden_states = hidden_states.reshape(batch_size, -1, self.head_dim)
+ dBx = (dB * hidden_states[..., None]).to(device=cache_device)
+
+ # State calculation
+ ssm_states = cache_params.layers[self.layer_idx].recurrent_states * dA + dBx
+ ssm_states = cache_params.update_recurrent_state(ssm_states, self.layer_idx)
+
+ # Subsequent output
+ # [bsz, n_groups * state_size] -> [bsz, num_heads, state_size]
+ C = C.reshape(batch_size, self.n_groups, -1)[..., None, :]
+ C = C.expand(batch_size, self.n_groups, self.num_heads // self.n_groups, C.shape[-1]).contiguous()
+ C = C.reshape(batch_size, -1, C.shape[-1])
+ # [bsz, num_heads, head_dim]
+
+ ssm_states = ssm_states.to(device=C.device, dtype=C.dtype) # Shape: [b, h, d, n]
+ # Reshape ssm_states to merge the first two dimensions
+ ssm_states_reshaped = ssm_states.view(batch_size * self.num_heads, self.head_dim, self.ssm_state_size) # Shape: [b*h, d, n]
+ C_reshaped = C.view(batch_size * self.num_heads, self.ssm_state_size, 1) # Shape: [b*h, n, 1]
+ y = torch.bmm(ssm_states_reshaped, C_reshaped)
+ y = y.view(batch_size, self.num_heads, self.head_dim)
+
+ # D skip connection
+ # [num_heads] -> [num_heads, head_dim]
+ D = self.D[..., None].expand(self.D.shape[0], self.head_dim)
+ y = (y + hidden_states * D).to(y.dtype)
+
+ # [bsz, num_heads, head_dim] -> [bsz, 1, intermediate_size]
+ y = y.reshape(batch_size, -1)[:, None, ...]
+ else:
+ # begin ssd naive implementation without einsums
+ dt = nn.functional.softplus(dt + self.dt_bias)
+ dt = torch.clamp(dt, self.time_step_limit[0], self.time_step_limit[1])
+ hidden_states = hidden_states.reshape(batch_size, seq_len, -1, self.head_dim).float()
+ B = B.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()
+ C = C.reshape(batch_size, seq_len, -1, self.ssm_state_size).float()
+ B = B.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
+ C = C.repeat_interleave(self.num_heads // self.n_groups, dim=2, output_size=self.num_heads)
+ pad_size = (self.chunk_size - seq_len % self.chunk_size) % self.chunk_size
+
+ D_residual = self.D[..., None] * pad_tensor_by_size(hidden_states, pad_size)
+
+ # Discretize x and A
+ hidden_states = hidden_states * dt[..., None]
+ A = A.to(hidden_states.dtype) * dt
+
+ # Rearrange into blocks/chunks
+ hidden_states, A, B, C = [reshape_into_chunks(t, pad_size, self.chunk_size) for t in (hidden_states, A, B, C)]
+
+ # [bsz, -1, chunk_size, num_heads] -> [bsz, num_heads, -1, chunk_size]
+ A = A.permute(0, 3, 1, 2)
+ A_cumsum = torch.cumsum(A, dim=-1)
+
+ # 1. Compute the output for each intra-chunk (diagonal blocks)
+ # This is the analog of a causal mask
+ L = torch.exp(segment_sum(A))
+
+ # Contraction of C and B to get G (attention-weights like)
+ G_intermediate = C[:, :, :, None, :, :] * B[:, :, None, :, :, :] # shape: (b, c, l, s, h, n)
+ G = G_intermediate.sum(dim=-1) # shape: (b, c, l, s, h)
+
+ # Compute M, equivalent to applying attention mask to weights
+ M_intermediate = G[..., None] * L.permute(0, 2, 3, 4, 1)[..., None]
+ M = M_intermediate.sum(dim=-1)
+
+ # Compute Y_diag (apply to values)
+ Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(dim=3)
+
+ # 2. Compute the state for each intra-chunk
+ # (right term of low-rank factorization of off-diagonal blocks; B terms)
+ decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum)
+ B_decay = B * decay_states.permute(0, -2, -1, 1)[..., None]
+ states = (B_decay[..., None, :] * hidden_states[..., None]).sum(dim=2)
+
+ # 3. Compute the inter-chunk SSM recurrence; produces correct SSM states at chunk boundaries
+ # (middle term of factorization of off-diag blocks; A terms)
+ previous_states = torch.zeros_like(states[:, :1])
+ states = torch.cat([previous_states, states], dim=1)
+ decay_chunk = torch.exp(segment_sum(nn.functional.pad(A_cumsum[:, :, :, -1], (1, 0))))
+ decay_chunk = decay_chunk.transpose(1, 3)
+ new_states = (decay_chunk[..., None, None] * states[:, :, None, ...]).sum(dim=1)
+ states, ssm_state = new_states[:, :-1], new_states[:, -1]
+
+ # 4. Compute state -> output conversion per chunk
+ # (left term of low-rank factorization of off-diagonal blocks; C terms)
+ state_decay_out = torch.exp(A_cumsum)
+ C_times_states = (C[..., None, :] * states[:, :, None, ...])
+ state_decay_out_permuted = state_decay_out.permute(0, 2, 3, 1)
+ Y_off = (C_times_states.sum(-1) * state_decay_out_permuted[..., None])
+
+ # Add output of intra-chunk and inter-chunk terms (diagonal and off-diagonal blocks)
+ y = Y_diag + Y_off
+ # [bsz, -1, self.chunk_size, num_heads, head_dim] -> [bsz, (padded) seq_len, num_heads, head_dim]
+ y = y.reshape(batch_size, -1, self.num_heads, self.head_dim)
+
+ y = y + D_residual
+ # Cutting off padded chunks
+ if pad_size > 0:
+ y = y[:, :seq_len, :, :]
+ y = y.reshape(batch_size, seq_len, -1)
+
+ # Init cache
+ if ssm_state is not None and cache_params is not None:
+ ssm_state = cache_params.update_recurrent_state(ssm_state, self.layer_idx)
+
+ if self.mamba_rms_norm:
+ scan_output = self.norm(y, gate)
+ else:
+ scan_output = y * torch.nn.functional.silu(gate)
+
+ # end ssd naive
+
+ # 4. Final linear projection
+ contextualized_states = self.out_proj(scan_output.to(dtype)) # [batch, seq_len, hidden_size]
+ return contextualized_states
+ # fmt: on
+
+ def forward(
+ self,
+ hidden_states,
+ cache_params: Cache | None = None,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs,
+ ):
+ if is_fast_path_available and "cuda" in self.in_proj.weight.device.type and not is_torchdynamo_compiling():
+ return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask)
+ dtype = hidden_states.dtype
+ if attention_mask is not None and attention_mask.shape[1] > 1 and attention_mask.shape[0] > 1:
+ # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66
+ hidden_states = (hidden_states * attention_mask[:, :, None]).to(dtype)
+
+ return self.torch_forward(hidden_states, cache_params, attention_mask)
+
+
+class FalconH1MLP(LlamaMLP):
+ def __init__(self, config: FalconH1Config):
+ super().__init__(config)
+ self.gate_multiplier, self.down_multiplier = config.mlp_multipliers
+
+ def forward(self, x):
+ y = self.up_proj(x) * self.act_fn(self.gate_proj(x) * self.gate_multiplier)
+ y = self.down_proj(y) * self.down_multiplier
+ return y
+
+
+class FalconH1RMSNorm(LlamaRMSNorm):
+ pass
+
+
+class FalconH1DecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: FalconH1Config, layer_idx: int):
+ super().__init__()
+ self.feed_forward = FalconH1MLP(config)
+
+ head_dim = config.hidden_size // config.num_attention_heads
+ self.channels_attn = config.num_attention_heads * head_dim + 2 * config.num_key_value_heads * head_dim
+
+ self.mamba = FalconH1Mixer(config=config, layer_idx=layer_idx)
+
+ self.self_attn = FalconH1Attention(config, layer_idx)
+
+ self.attention_in_multiplier = config.attention_in_multiplier
+ self.ssm_out_multiplier = config.ssm_out_multiplier
+ self.attn_out_multiplier = config.attention_out_multiplier
+
+ self.input_layernorm = FalconH1RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.pre_ff_layernorm = FalconH1RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ mamba_attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs,
+ ) -> tuple[torch.FloatTensor]:
+ """
+ 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.
+ past_key_values (`Cache`, *optional*): cached past key and value projection states
+ 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`).
+ position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
+ Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
+ with `head_dim` being the embedding dimension of each attention head.
+ kwargs (`dict`, *optional*):
+ Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
+ into the model
+ """
+
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ mamba_hidden_states = self.mamba(
+ hidden_states=hidden_states,
+ cache_params=past_key_values,
+ attention_mask=mamba_attention_mask,
+ )
+ mamba_hidden_states = mamba_hidden_states * self.ssm_out_multiplier
+
+ attention_hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states * self.attention_in_multiplier,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ attention_hidden_states = attention_hidden_states * self.attn_out_multiplier
+
+ hidden_states = mamba_hidden_states + attention_hidden_states
+
+ # residual connection after attention
+ hidden_states = residual + hidden_states
+
+ # feed-forward
+ residual = hidden_states
+ hidden_states = self.pre_ff_layernorm(hidden_states)
+ hidden_states = self.feed_forward(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return (hidden_states,)
+
+
+@auto_docstring
+class FalconH1PreTrainedModel(PreTrainedModel):
+ config: FalconH1Config
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["FalconH1DecoderLayer"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _is_stateful = True
+
+ _can_record_outputs = {
+ "hidden_states": FalconH1DecoderLayer,
+ "attentions": FalconH1Attention,
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, FalconH1Mixer):
+ init.ones_(module.dt_bias)
+ init.copy_(module.A_log, torch.log(torch.arange(1, module.num_heads + 1)))
+ init.ones_(module.D)
+ elif isinstance(module, FalconH1Model):
+ mup_vector = compute_mup_vector(module.config)
+ for layer in module.layers:
+ init.copy_(layer.mamba.mup_vector, mup_vector)
+
+
+def compute_mup_vector(config):
+ """
+ Computes the MuP vector based on model configuration.
+
+ FalconH1 applies different MuP multiplier for each dimension of the hidden states.
+ The MuP vector is partitioned into chunks, and each chunk is multiplied with its
+ corresponding projected dimension.
+
+ Args:
+ config: FalconH1Config object
+
+ Returns:
+ torch.Tensor: The computed MuP vector
+ """
+ # We'll need some values from the config to compute the vector dimensions
+ intermediate_size = (
+ config.mamba_d_ssm if config.mamba_d_ssm is not None else int(config.mamba_expand * config.hidden_size)
+ )
+ groups_time_state_size = config.mamba_n_groups * config.mamba_d_state
+ num_heads = config.mamba_n_heads
+ zxbcdt_multipliers = config.ssm_multipliers
+
+ vector_shape = 2 * intermediate_size + 2 * groups_time_state_size + num_heads
+ mup_vector = torch.ones(1, 1, vector_shape)
+
+ # Apply multipliers to different sections of the vector
+ mup_vector[:, :, :intermediate_size] *= zxbcdt_multipliers[0]
+ mup_vector[:, :, intermediate_size : 2 * intermediate_size] *= zxbcdt_multipliers[1]
+ mup_vector[:, :, 2 * intermediate_size : 2 * intermediate_size + groups_time_state_size] *= zxbcdt_multipliers[2]
+ mup_vector[
+ :, :, 2 * intermediate_size + groups_time_state_size : 2 * intermediate_size + 2 * groups_time_state_size
+ ] *= zxbcdt_multipliers[3]
+ mup_vector[:, :, 2 * intermediate_size + 2 * groups_time_state_size :] *= zxbcdt_multipliers[4]
+
+ return mup_vector
+
+
+@auto_docstring
+# Adapted from transformers.models.jamba.modeling_jamba.JambaModel
+class FalconH1Model(FalconH1PreTrainedModel):
+ def __init__(self, config: FalconH1Config):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ decoder_layers = []
+ for i in range(config.num_hidden_layers):
+ decoder_layers.append(FalconH1DecoderLayer(config, layer_idx=i))
+ self.layers = nn.ModuleList(decoder_layers)
+
+ self._attn_implementation = config._attn_implementation
+ self.final_layernorm = FalconH1RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = FalconH1RotaryEmbedding(config=config)
+
+ self.embedding_multiplier = config.embedding_multiplier
+ self.lm_head_multiplier = config.lm_head_multiplier
+
+ self.gradient_checkpointing = False
+ # Compute the MuP vector once and register it for all layers
+ mup_vector = compute_mup_vector(config)
+ for layer in self.layers:
+ layer.mamba.register_buffer("mup_vector", mup_vector.clone(), persistent=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids) * self.embedding_multiplier
+ hidden_states = inputs_embeds
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(hidden_states.shape[1], device=hidden_states.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+ mamba_mask = self._update_mamba_mask(attention_mask, past_key_values)
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for decoder_layer in self.layers:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ mamba_attention_mask=mamba_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ )
+
+ hidden_states = layer_outputs[0]
+
+ hidden_states = self.final_layernorm(hidden_states)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+ def _update_mamba_mask(self, attention_mask, past_key_values):
+ """
+ No need for zeroing states when
+ 1. Cached forward
+ 2. Attending to all inputs
+ """
+ mamba_mask = attention_mask
+ if (past_key_values is not None and past_key_values.has_previous_state()) or (
+ attention_mask is not None and torch.all(attention_mask == 1)
+ ):
+ mamba_mask = None
+ return mamba_mask
+
+
+class FalconH1ForCausalLM(LlamaForCausalLM):
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs,
+ ) -> tuple | CausalLMOutputWithPast:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, FalconH1ForCausalLM
+
+ >>> model = FalconH1ForCausalLM.from_pretrained("...")
+ >>> tokenizer = AutoTokenizer.from_pretrained("...")
+
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+ ```"""
+ 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,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :]) * self.model.lm_head_multiplier
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ position_ids=None,
+ use_cache=True,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ kwargs["logits_to_keep"] = self.config.num_logits_to_keep
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ position_ids=position_ids,
+ use_cache=use_cache,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ return model_inputs
+
+
+__all__ = ["FalconH1Model", "FalconH1ForCausalLM", "FalconH1PreTrainedModel"]
diff --git a/third_party/transformers/src/transformers/models/fast_vlm/__init__.py b/third_party/transformers/src/transformers/models/fast_vlm/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..949f087650dd99ea1b04cb052708b9123b08b557
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/fast_vlm/__init__.py
@@ -0,0 +1,27 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_fast_vlm import *
+ from .modeling_fast_vlm import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/fast_vlm/configuration_fast_vlm.py b/third_party/transformers/src/transformers/models/fast_vlm/configuration_fast_vlm.py
new file mode 100644
index 0000000000000000000000000000000000000000..25c61da6507f248a8594cc0d10c1354358d6be31
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/fast_vlm/configuration_fast_vlm.py
@@ -0,0 +1,111 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/fast_vlm/modular_fast_vlm.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_fast_vlm.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+@auto_docstring(checkpoint="KamilaMila/FastVLM-7B")
+@strict
+class FastVlmConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import FastVlmForConditionalGeneration, FastVlmConfig
+
+ >>> # Initializing a FastVLM-7B style configuration
+ >>> configuration = FastVlmConfig()
+
+ >>> # Initializing a model from the FastVLM-7B style configuration
+ >>> model = FastVlmForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "fast_vlm"
+ attribute_map = {
+ "image_token_id": "image_token_index",
+ }
+ sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig}
+
+ vision_config: dict | PreTrainedConfig | None = None
+ text_config: dict | PreTrainedConfig | None = None
+ image_token_index: int = 151646
+ image_seq_length: int = 576
+ projector_hidden_act: str = "gelu"
+ vision_feature_select_strategy: str = "full"
+ vision_feature_layer: int | list[int] = -1
+ multimodal_projector_bias: bool = True
+ tie_word_embeddings: bool = False
+
+ def __post_init__(self, **kwargs):
+ if isinstance(self.vision_config, dict):
+ self.vision_config["model_type"] = self.vision_config.get("model_type", "timm_wrapper")
+ self.vision_config = CONFIG_MAPPING[self.vision_config["model_type"]](**self.vision_config)
+ elif self.vision_config is None:
+ self.vision_config = CONFIG_MAPPING["timm_wrapper"](
+ architecture="fastvit_mci3",
+ do_pooling=True,
+ global_pool="avg",
+ hidden_size=3072,
+ initializer_range=0.02,
+ model_args={"inference_mode": True},
+ )
+
+ if isinstance(self.text_config, dict):
+ self.text_config["model_type"] = self.text_config.get("model_type", "qwen2")
+ self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config)
+ elif self.text_config is None:
+ self.text_config = CONFIG_MAPPING["qwen2"](
+ hidden_size=3584,
+ vocab_size=152128,
+ intermediate_size=18944,
+ num_attention_heads=28,
+ num_key_value_heads=4,
+ num_hidden_layers=28,
+ )
+ # The default value is `False` but this config is used with many model types
+ # Attr `tie_word_embeddings` was saved in text config for those models, so we
+ # need an ugly workaround and forward-pass the attr from text config
+ if not self.tie_word_embeddings and self.text_config.tie_word_embeddings:
+ self.tie_word_embeddings = self.text_config.tie_word_embeddings
+
+ super().__post_init__(**kwargs)
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ if self.vision_feature_select_strategy != "full":
+ raise ValueError(
+ f"Unexpected select feature strategy: {self.vision_feature_select_strategy}. Only 'full' is supported in FastVLM."
+ )
+
+ if self.vision_feature_layer != -1:
+ raise ValueError(
+ f"Unexpected vision feature layer: {self.vision_feature_layer}. Only -1 is supported in FastVLM."
+ )
+
+
+__all__ = ["FastVlmConfig"]
diff --git a/third_party/transformers/src/transformers/models/fast_vlm/convert_fastvlm_weights_to_hf.py b/third_party/transformers/src/transformers/models/fast_vlm/convert_fastvlm_weights_to_hf.py
new file mode 100644
index 0000000000000000000000000000000000000000..eab68749bd30cdf8361cf94db6bb56c964fcb92e
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/fast_vlm/convert_fastvlm_weights_to_hf.py
@@ -0,0 +1,249 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+import argparse
+import glob
+import os
+import re
+from io import BytesIO
+
+import httpx
+import torch
+from huggingface_hub import snapshot_download
+from PIL import Image
+from safetensors import safe_open
+
+from transformers import (
+ AddedToken,
+ AutoConfig,
+ AutoTokenizer,
+ CLIPImageProcessor,
+ FastVlmConfig,
+ FastVlmForConditionalGeneration,
+ LlavaProcessor,
+)
+
+
+os.environ["TIMM_FUSED_ATTN"] = (
+ "0" # to avoid logits diverging, needed because the original implementation uses regular (not fused) atteniton
+)
+
+KEYS_TO_MODIFY_MAPPING = {
+ "model.vision_tower.vision_tower.model": "model.vision_tower.timm_model",
+ "patch_embed": "stem",
+ "layers": "language_model.layers",
+ "embed_tokens": "language_model.embed_tokens",
+ "layer_scale_1": "layer_scale_1.gamma",
+ "layer_scale_2": "layer_scale_2.gamma",
+ "mm_projector.0": "multi_modal_projector.linear_1",
+ "mm_projector.2": "multi_modal_projector.linear_2",
+ "conv_exp": "final_conv",
+ "se.reduce": "se.fc1",
+ "se.expand": "se.fc2",
+ "convffn": "mlp",
+ "lkb_reparam": "reparam_conv",
+}
+
+
+def map_to_stage(number):
+ number = int(number)
+ if number == 0:
+ return 0
+ if number in {1, 2}:
+ return 1
+ if number in {3, 4}:
+ return 2
+ if number in {5, 6, 7}:
+ return 3
+ if number in {8, 9, 10}:
+ return 4
+
+
+def load_original_state_dict(model_id):
+ directory_path = snapshot_download(repo_id=model_id, allow_patterns=["*.safetensors"])
+
+ original_state_dict = {}
+ for path in glob.glob(f"{directory_path}/*"):
+ if path.endswith(".safetensors"):
+ with safe_open(path, framework="pt", device="cpu") as f:
+ for key in f.keys():
+ original_state_dict[key] = f.get_tensor(key)
+
+ if "model.vision_tower.vision_tower.model.head.proj" in original_state_dict:
+ del original_state_dict["model.vision_tower.vision_tower.model.head.proj"]
+ return original_state_dict
+
+
+def convert_state_dict_to_hf(state_dict):
+ new_state_dict = {}
+
+ single_pattern = r"network\.(\d{1,2})"
+ double_pattern = r"network\.(\d{1,2})\.(\d{1,2})"
+ pos_embedding_pattern = r"stages\.(\d{1,2})\.reparam_conv"
+
+ for key, value in state_dict.items():
+ if key.endswith("layer_scale"):
+ key = key.replace("layer_scale", "layer_scale.gamma")
+ if key.startswith("model.norm"):
+ key = key.replace("model.norm", "model.language_model.norm")
+ if "token_mixer" not in key:
+ key = key.replace(".proj.", ".downsample.proj.")
+
+ matches = re.findall(double_pattern, key)
+ if len(matches) == 1:
+ match = matches[0]
+ key = key.replace(f"network.{match[0]}.{match[1]}", f"stages.{map_to_stage(match[0])}.blocks.{match[1]}")
+
+ matches = re.findall(single_pattern, key)
+ if len(matches) == 1:
+ match = matches[0]
+ key = key.replace(f"network.{match[0]}", f"stages.{map_to_stage(match[0])}")
+
+ matches = re.findall(pos_embedding_pattern, key)
+ if len(matches) == 1:
+ match = matches[0]
+ key = key.replace(f"stages.{match[0]}", f"stages.{match[0]}.pos_emb")
+
+ for key_to_modify, new_key in KEYS_TO_MODIFY_MAPPING.items():
+ if key_to_modify in key:
+ key = key.replace(key_to_modify, new_key)
+
+ new_state_dict[key] = value
+ return new_state_dict
+
+
+def convert_fastvlm_to_hf(text_model_id, vision_model_id, output_hub_path, old_state_dict_id):
+ torch.set_default_dtype(torch.bfloat16)
+
+ text_config = AutoConfig.from_pretrained(text_model_id)
+ vision_config = AutoConfig.from_pretrained(vision_model_id)
+ vision_config.model_args = {"inference_mode": True}
+ vision_config.hidden_size = vision_config.num_features
+ vision_config.label2id = {}
+ vision_config.id2label = {}
+ config = FastVlmConfig(
+ text_config=text_config,
+ vision_config=vision_config,
+ )
+ config.vision_feature_select_strategy = "full"
+ config.vision_feature_layer = -1
+ config.image_token_index = 151646
+ config.image_seq_length = 256
+
+ tokenizer = AutoTokenizer.from_pretrained(
+ text_model_id,
+ chat_template="{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}{% endif %}{{'<|im_start|>' + message['role'] + '\n'}}{# Render all images first #}{% for content in message['content'] | selectattr('type', 'equalto', 'image') %}{{ '' }}{% endfor %}{# Render all text next #}{% for content in message['content'] | selectattr('type', 'equalto', 'text') %}{{ '\n' + content['text'] }}{% endfor %}{{'<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
+ )
+
+ tokenizer.add_tokens(AddedToken("", special=True, normalized=False), special_tokens=True)
+ image_processor = CLIPImageProcessor(
+ crop_size={"height": 1024, "width": 1024},
+ image_mean=[0.0, 0.0, 0.0],
+ image_std=[1.0, 1.0, 1.0],
+ size={"shortest_edge": 1024},
+ )
+
+ processor = LlavaProcessor(tokenizer=tokenizer, image_processor=image_processor)
+ processor.patch_size = 64 # effective patch size (2^6)
+
+ model = FastVlmForConditionalGeneration(config)
+
+ state_dict = load_original_state_dict(old_state_dict_id)
+ state_dict = convert_state_dict_to_hf(state_dict)
+ model.load_state_dict(state_dict, strict=True, assign=True)
+
+ pre_expansion_embeddings = model.language_model.embed_tokens.weight.data
+ mu = torch.mean(pre_expansion_embeddings, dim=0).float()
+ n = pre_expansion_embeddings.size()[0]
+ sigma = ((pre_expansion_embeddings - mu).T @ (pre_expansion_embeddings - mu)) / n
+ dist = torch.distributions.multivariate_normal.MultivariateNormal(mu, covariance_matrix=1e-5 * sigma)
+
+ # We add an image token so we resize the model and pad to 64 for performance reasons
+ pad_shape = 64
+ vocab_size = config.text_config.vocab_size
+ model.resize_token_embeddings(config.text_config.vocab_size + 1, pad_shape)
+ model.language_model.embed_tokens.weight.data[vocab_size:] = torch.stack(
+ tuple(dist.sample() for _ in range(model.language_model.embed_tokens.weight.data[vocab_size:].shape[0])),
+ dim=0,
+ )
+ model.lm_head.weight.data[vocab_size:] = torch.stack(
+ tuple(dist.sample() for _ in range(model.lm_head.weight.data[vocab_size:].shape[0])),
+ dim=0,
+ )
+
+ conversation = [{"role": "user", "content": [{"type": "text", "text": "What are these?"}, {"type": "image"}]}]
+ prompt = tokenizer.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
+
+ url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ with httpx.stream("GET", url) as response:
+ raw_image = Image.open(BytesIO(response.read()))
+ inputs = processor(images=raw_image, text=prompt, return_tensors="pt").to("cuda")
+ inputs = {k: (v.to(torch.bfloat16) if v.dtype == torch.float32 else v) for k, v in inputs.items()}
+
+ model = model.cuda()
+ model.eval()
+ with torch.no_grad():
+ logits = model(**inputs).logits
+
+ # in order to get the same logits as in the Apple repo, we need to manually replace the original (Apple) LayerNorm2D with Timm's LayerNorm2D or vice versa
+ # otherwise numerical errors accumulate
+ if output_hub_path == "KamilaMila/FastVLM-0.5B":
+ expected_shape = torch.Size([1, 280, 152000])
+ expected_slice = torch.tensor([4.1250, 9.6875, 11.1875], device="cuda")
+ elif output_hub_path == "KamilaMila/FastVLM-1.5B":
+ expected_shape = torch.Size([1, 280, 152000])
+ expected_slice = torch.tensor([3.3750, 11.5000, 11.8125], device="cuda")
+ elif output_hub_path == "KamilaMila/FastVLM-7B":
+ expected_shape = torch.Size([1, 280, 152128])
+ expected_slice = torch.tensor([3.8281, 9.0625, 7.9062], device="cuda")
+
+ logits_slice = logits[0, -1, :3]
+ assert torch.allclose(expected_slice, logits_slice, atol=1e-8)
+ assert logits.shape == expected_shape
+
+ model.push_to_hub(output_hub_path)
+ processor.push_to_hub(output_hub_path)
+ print("Successfully pushed to hub!")
+
+
+def main():
+ parser = argparse.ArgumentParser(
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ )
+
+ parser.add_argument(
+ "--text_model_id",
+ default="Qwen/Qwen2-1.5B",
+ help="Hub location of the text model",
+ )
+ parser.add_argument(
+ "--vision_model_id",
+ default="timm/fastvit_mci3.apple_mclip2_dfndr2b",
+ help="Hub location of the vision model",
+ )
+ parser.add_argument(
+ "--output_hub_path",
+ default="KamilaMila/FastVLM-1.5B",
+ help="Location on the hub of the converted model",
+ )
+ parser.add_argument(
+ "--old_state_dict_id",
+ default="apple/FastVLM-1.5B",
+ help="Location on the hub of the raw state dict of the original model.",
+ )
+ args = parser.parse_args()
+ convert_fastvlm_to_hf(args.text_model_id, args.vision_model_id, args.output_hub_path, args.old_state_dict_id)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/third_party/transformers/src/transformers/models/fast_vlm/modeling_fast_vlm.py b/third_party/transformers/src/transformers/models/fast_vlm/modeling_fast_vlm.py
new file mode 100644
index 0000000000000000000000000000000000000000..85c2eeb82b6488882d23b974d0c3937a9ff393d6
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/fast_vlm/modeling_fast_vlm.py
@@ -0,0 +1,425 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/fast_vlm/modular_fast_vlm.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_fast_vlm.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache
+from ...generation import GenerationMixin
+from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, torch_compilable_check
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ..auto import AutoModel
+from .configuration_fast_vlm import FastVlmConfig
+
+
+class FastVlmMultiModalProjector(nn.Module):
+ def __init__(self, config: FastVlmConfig):
+ super().__init__()
+ self.linear_1 = nn.Linear(
+ config.vision_config.hidden_size,
+ config.text_config.hidden_size,
+ bias=config.multimodal_projector_bias,
+ )
+ self.act = ACT2FN[config.projector_hidden_act]
+ self.linear_2 = nn.Linear(
+ config.text_config.hidden_size, config.text_config.hidden_size, bias=config.multimodal_projector_bias
+ )
+
+ def forward(self, image_features):
+ hidden_states = self.linear_1(image_features)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.linear_2(hidden_states)
+ return hidden_states
+
+
+@auto_docstring
+class FastVlmPreTrainedModel(PreTrainedModel):
+ config: FastVlmConfig
+ base_model_prefix = "model"
+ input_modalities = ("image", "text")
+ supports_gradient_checkpointing = True
+ _skip_keys_device_placement = "past_key_values"
+
+ _supports_flash_attn = True
+ _supports_sdpa = True
+
+ _can_compile_fullgraph = True
+ _supports_flex_attn = True
+ _supports_attention_backend = True
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for FastVlm outputs, with hidden states and attentions.
+ """
+)
+class FastVlmModelOutputWithPast(BaseModelOutputWithPast):
+ r"""
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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.
+ image_hidden_states (`torch.FloatTensor`, *optional*):
+ A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+ image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+ """
+
+ image_hidden_states: torch.FloatTensor | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ The FastVlm model which consists of a vision backbone and a language model, without a language modeling head.
+ """
+)
+class FastVlmModel(FastVlmPreTrainedModel):
+ def __init__(self, config: FastVlmConfig):
+ super().__init__(config)
+ self.vision_tower = AutoModel.from_config(config.vision_config)
+
+ self.multi_modal_projector = FastVlmMultiModalProjector(config)
+ self.language_model = AutoModel.from_config(config.text_config)
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.language_model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.language_model.set_input_embeddings(value)
+
+ @merge_with_config_defaults
+ @can_return_tuple
+ @auto_docstring(
+ custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection."
+ )
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ vision_feature_layer: int | list[int] | list[int] | None = None,
+ vision_feature_select_strategy: str | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`):
+ The tensors corresponding to the input images.
+ vision_feature_layer (`Union[int, list[int]]`, *optional*):
+ The index/indices of the layer to select the vision feature. Only -1 supported.
+ vision_feature_select_strategy (`str`, *optional*):
+ The feature selection strategy used to select the vision feature from the vision backbone.
+ Only "full" supported.
+ """
+ kwargs = {k: v for k, v in kwargs.items() if v is not None}
+ image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs)
+
+ # since the vision tower is hybrid in FastVLM, its output needs to be handled differently from Llava
+ selected_image_feature = image_outputs.last_hidden_state
+ selected_image_feature = selected_image_feature.flatten(2).permute(0, 2, 1)
+ image_features = self.multi_modal_projector(selected_image_feature)
+ image_outputs.pooler_output = list(image_features)
+
+ return image_outputs
+
+ def get_placeholder_mask(
+ self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor, image_features: torch.FloatTensor
+ ):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+ equal to the length of multimodal features. If the lengths are different, an error is raised.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ else:
+ special_image_mask = input_ids == self.config.image_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ n_image_features = image_features.shape[0] * image_features.shape[1]
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ torch_compilable_check(
+ inputs_embeds[special_image_mask].numel() == image_features.numel(),
+ f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {n_image_features}",
+ )
+ return special_image_mask
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ vision_feature_layer: int | list[int] | list[int] | None = None,
+ vision_feature_select_strategy: str | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | FastVlmModelOutputWithPast:
+ r"""
+ vision_feature_layer (`Union[int, list[int], NoneType]`, *optional*):
+ The index of the layer to select the vision feature. If multiple indices are provided, the vision feature of the
+ corresponding indices will be concatenated to form the vision features. Only -1 supported.
+ vision_feature_select_strategy (`str`, *optional*):
+ The feature selection strategy used to select the vision feature from the vision backbone. Only "full" supported.
+ """
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_features = self.get_image_features(
+ pixel_values=pixel_values,
+ vision_feature_layer=vision_feature_layer,
+ vision_feature_select_strategy=vision_feature_select_strategy,
+ return_dict=True,
+ ).pooler_output
+ image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ special_image_mask = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_features
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+ outputs = self.language_model(
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ return FastVlmModelOutputWithPast(
+ last_hidden_state=outputs.last_hidden_state,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ image_hidden_states=image_features if pixel_values is not None else None,
+ )
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for FastVlm causal language model (or autoregressive) outputs.
+ """
+)
+class FastVlmCausalLMOutputWithPast(ModelOutput):
+ r"""
+ 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 (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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.
+ image_hidden_states (`torch.FloatTensor`, *optional*):
+ A `torch.FloatTensor` of size `(batch_size, num_images, sequence_length, hidden_size)`.
+ image_hidden_states of the model produced by the vision encoder and after projecting the last hidden state.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ image_hidden_states: torch.FloatTensor | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ The FastVlm model which consists of a vision backbone and a language model.
+ """
+)
+class FastVlmForConditionalGeneration(FastVlmPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+
+ def __init__(self, config: FastVlmConfig):
+ super().__init__(config)
+ self.model = FastVlmModel(config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.model.set_input_embeddings(value)
+
+ def get_output_embeddings(self) -> nn.Module:
+ return self.lm_head
+
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ vision_feature_layer: int | list[int] | list[int] | None = None,
+ vision_feature_select_strategy: str | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ return self.model.get_image_features(
+ pixel_values=pixel_values,
+ vision_feature_layer=vision_feature_layer,
+ vision_feature_select_strategy=vision_feature_select_strategy,
+ **kwargs,
+ )
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ vision_feature_layer: int | list[int] | list[int] | None = None,
+ vision_feature_select_strategy: str | None = None,
+ labels: torch.LongTensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | FastVlmCausalLMOutputWithPast:
+ r"""
+ vision_feature_layer (`Union[int, list[int], NoneType]`, *optional*):
+ The index of the layer to select the vision feature. If multiple indices are provided, the vision feature of the
+ corresponding indices will be concatenated to form the vision features. Only -1 supported.
+ vision_feature_select_strategy (`str`, *optional*):
+ The feature selection strategy used to select the vision feature from the vision backbone. Only "full" supported.
+ 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]`.
+
+ Example:
+
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoProcessor, AutoModelForImageTextToText
+ >>> import torch
+
+ >>> device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ >>> model = AutoModelForImageTextToText.from_pretrained("KamilaMila/FastVLM-0.5B").to(device)
+ >>> processor = AutoProcessor.from_pretrained("KamilaMila/FastVLM-0.5B")
+
+ >>> conversation = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "What are these?"},
+ {"type": "image"}
+ ]
+ }
+ ]
+
+ >>> prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+ >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> inputs = processor(images=image, text=prompt, return_tensors="pt").to(device)
+
+ >>> # Generate
+ >>> generated_ids = model.generate(**inputs, max_new_tokens=15)
+ >>> print(processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0])
+ system\n You are a helpful assistant.\n user\n What are these?\n assistant\n The image depicts a traditional Chinese street...
+ ```"""
+ outputs = self.model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ vision_feature_layer=vision_feature_layer,
+ vision_feature_select_strategy=vision_feature_select_strategy,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+ )
+
+ return FastVlmCausalLMOutputWithPast(
+ 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_key_values=None,
+ inputs_embeds=None,
+ pixel_values=None,
+ attention_mask=None,
+ logits_to_keep=None,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ logits_to_keep=logits_to_keep,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ if is_first_iteration or not kwargs.get("use_cache", True):
+ # Pixel values are used only in the first iteration if available
+ # In subsequent iterations, they are already merged with text and cached
+ # NOTE: first iteration doesn't have to be prefill, it can be the first
+ # iteration with a question and cached system prompt (continue generate from cache)
+ model_inputs["pixel_values"] = pixel_values
+
+ return model_inputs
+
+
+__all__ = ["FastVlmForConditionalGeneration", "FastVlmModel", "FastVlmPreTrainedModel"]
diff --git a/third_party/transformers/src/transformers/models/fast_vlm/modular_fast_vlm.py b/third_party/transformers/src/transformers/models/fast_vlm/modular_fast_vlm.py
new file mode 100644
index 0000000000000000000000000000000000000000..d9db87ab27d4318d9fd6dace7b15ad4cb73ab732
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/fast_vlm/modular_fast_vlm.py
@@ -0,0 +1,337 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+
+import torch
+from huggingface_hub.dataclasses import strict
+from torch import nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_outputs import BaseModelOutputWithPooling
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ..auto import CONFIG_MAPPING
+from ..llava.configuration_llava import LlavaConfig
+from ..llava.modeling_llava import (
+ LlavaCausalLMOutputWithPast,
+ LlavaForConditionalGeneration,
+ LlavaModel,
+ LlavaModelOutputWithPast,
+ LlavaMultiModalProjector,
+ LlavaPreTrainedModel,
+)
+
+
+@auto_docstring(checkpoint="KamilaMila/FastVLM-7B")
+@strict
+class FastVlmConfig(LlavaConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import FastVlmForConditionalGeneration, FastVlmConfig
+
+ >>> # Initializing a FastVLM-7B style configuration
+ >>> configuration = FastVlmConfig()
+
+ >>> # Initializing a model from the FastVLM-7B style configuration
+ >>> model = FastVlmForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "fast_vlm"
+
+ vision_config: dict | PreTrainedConfig | None = None
+ text_config: dict | PreTrainedConfig | None = None
+ image_token_index: int = 151646
+ projector_hidden_act: str = "gelu"
+ vision_feature_select_strategy: str = "full"
+ vision_feature_layer: int | list[int] = -1
+ multimodal_projector_bias: bool = True
+ tie_word_embeddings: bool = False
+
+ def __post_init__(self, **kwargs):
+ if isinstance(self.vision_config, dict):
+ self.vision_config["model_type"] = self.vision_config.get("model_type", "timm_wrapper")
+ self.vision_config = CONFIG_MAPPING[self.vision_config["model_type"]](**self.vision_config)
+ elif self.vision_config is None:
+ self.vision_config = CONFIG_MAPPING["timm_wrapper"](
+ architecture="fastvit_mci3",
+ do_pooling=True,
+ global_pool="avg",
+ hidden_size=3072,
+ initializer_range=0.02,
+ model_args={"inference_mode": True},
+ )
+
+ if isinstance(self.text_config, dict):
+ self.text_config["model_type"] = self.text_config.get("model_type", "qwen2")
+ self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config)
+ elif self.text_config is None:
+ self.text_config = CONFIG_MAPPING["qwen2"](
+ hidden_size=3584,
+ vocab_size=152128,
+ intermediate_size=18944,
+ num_attention_heads=28,
+ num_key_value_heads=4,
+ num_hidden_layers=28,
+ )
+ # The default value is `False` but this config is used with many model types
+ # Attr `tie_word_embeddings` was saved in text config for those models, so we
+ # need an ugly workaround and forward-pass the attr from text config
+ if not self.tie_word_embeddings and self.text_config.tie_word_embeddings:
+ self.tie_word_embeddings = self.text_config.tie_word_embeddings
+
+ PreTrainedConfig.__post_init__(**kwargs)
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ if self.vision_feature_select_strategy != "full":
+ raise ValueError(
+ f"Unexpected select feature strategy: {self.vision_feature_select_strategy}. Only 'full' is supported in FastVLM."
+ )
+
+ if self.vision_feature_layer != -1:
+ raise ValueError(
+ f"Unexpected vision feature layer: {self.vision_feature_layer}. Only -1 is supported in FastVLM."
+ )
+
+
+class FastVlmMultiModalProjector(LlavaMultiModalProjector):
+ def __init__(self, config: FastVlmConfig):
+ nn.Module.__init__()
+ self.linear_1 = nn.Linear(
+ config.vision_config.hidden_size,
+ config.text_config.hidden_size,
+ bias=config.multimodal_projector_bias,
+ )
+ self.act = ACT2FN[config.projector_hidden_act]
+ self.linear_2 = nn.Linear(
+ config.text_config.hidden_size, config.text_config.hidden_size, bias=config.multimodal_projector_bias
+ )
+
+
+class FastVlmPreTrainedModel(LlavaPreTrainedModel):
+ pass
+
+
+class FastVlmModelOutputWithPast(LlavaModelOutputWithPast):
+ pass
+
+
+class FastVlmModel(LlavaModel):
+ def __init__(self, config: FastVlmConfig):
+ super().__init__(config)
+
+ @merge_with_config_defaults
+ @can_return_tuple
+ @auto_docstring(
+ custom_intro="Obtains image last hidden states from the vision tower and apply multimodal projection."
+ )
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ vision_feature_layer: int | list[int] | list[int] | None = None,
+ vision_feature_select_strategy: str | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`):
+ The tensors corresponding to the input images.
+ vision_feature_layer (`Union[int, list[int]]`, *optional*):
+ The index/indices of the layer to select the vision feature. Only -1 supported.
+ vision_feature_select_strategy (`str`, *optional*):
+ The feature selection strategy used to select the vision feature from the vision backbone.
+ Only "full" supported.
+ """
+ kwargs = {k: v for k, v in kwargs.items() if v is not None}
+ image_outputs = self.vision_tower(pixel_values, return_dict=True, **kwargs)
+
+ # since the vision tower is hybrid in FastVLM, its output needs to be handled differently from Llava
+ selected_image_feature = image_outputs.last_hidden_state
+ selected_image_feature = selected_image_feature.flatten(2).permute(0, 2, 1)
+ image_features = self.multi_modal_projector(selected_image_feature)
+ image_outputs.pooler_output = list(image_features)
+
+ return image_outputs
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ vision_feature_layer: int | list[int] | list[int] | None = None,
+ vision_feature_select_strategy: str | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | FastVlmModelOutputWithPast:
+ r"""
+ vision_feature_layer (`Union[int, list[int], NoneType]`, *optional*):
+ The index of the layer to select the vision feature. If multiple indices are provided, the vision feature of the
+ corresponding indices will be concatenated to form the vision features. Only -1 supported.
+ vision_feature_select_strategy (`str`, *optional*):
+ The feature selection strategy used to select the vision feature from the vision backbone. Only "full" supported.
+ """
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_features = self.get_image_features(
+ pixel_values=pixel_values,
+ vision_feature_layer=vision_feature_layer,
+ vision_feature_select_strategy=vision_feature_select_strategy,
+ return_dict=True,
+ ).pooler_output
+ image_features = torch.cat(image_features, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ special_image_mask = self.get_placeholder_mask(
+ input_ids, inputs_embeds=inputs_embeds, image_features=image_features
+ )
+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, image_features)
+
+ outputs = self.language_model(
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ return FastVlmModelOutputWithPast(
+ last_hidden_state=outputs.last_hidden_state,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ image_hidden_states=image_features if pixel_values is not None else None,
+ )
+
+
+class FastVlmCausalLMOutputWithPast(LlavaCausalLMOutputWithPast):
+ pass
+
+
+@auto_docstring(
+ custom_intro="""
+ The FastVlm model which consists of a vision backbone and a language model.
+ """
+)
+class FastVlmForConditionalGeneration(LlavaForConditionalGeneration):
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ vision_feature_layer: int | list[int] | list[int] | None = None,
+ vision_feature_select_strategy: str | None = None,
+ labels: torch.LongTensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | FastVlmCausalLMOutputWithPast:
+ r"""
+ vision_feature_layer (`Union[int, list[int], NoneType]`, *optional*):
+ The index of the layer to select the vision feature. If multiple indices are provided, the vision feature of the
+ corresponding indices will be concatenated to form the vision features. Only -1 supported.
+ vision_feature_select_strategy (`str`, *optional*):
+ The feature selection strategy used to select the vision feature from the vision backbone. Only "full" supported.
+ 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]`.
+
+ Example:
+
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoProcessor, AutoModelForImageTextToText
+ >>> import torch
+
+ >>> device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ >>> model = AutoModelForImageTextToText.from_pretrained("KamilaMila/FastVLM-0.5B").to(device)
+ >>> processor = AutoProcessor.from_pretrained("KamilaMila/FastVLM-0.5B")
+
+ >>> conversation = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "text", "text": "What are these?"},
+ {"type": "image"}
+ ]
+ }
+ ]
+
+ >>> prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+ >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> inputs = processor(images=image, text=prompt, return_tensors="pt").to(device)
+
+ >>> # Generate
+ >>> generated_ids = model.generate(**inputs, max_new_tokens=15)
+ >>> print(processor.batch_decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0])
+ system\n You are a helpful assistant.\n user\n What are these?\n assistant\n The image depicts a traditional Chinese street...
+ ```"""
+ outputs = self.model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ vision_feature_layer=vision_feature_layer,
+ vision_feature_select_strategy=vision_feature_select_strategy,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+ )
+
+ return FastVlmCausalLMOutputWithPast(
+ 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,
+ )
+
+
+__all__ = ["FastVlmForConditionalGeneration", "FastVlmModel", "FastVlmPreTrainedModel", "FastVlmConfig"]
diff --git a/third_party/transformers/src/transformers/models/flava/__init__.py b/third_party/transformers/src/transformers/models/flava/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..06b27a7e05f3e58dd0ac66bc4476621ebda2d5bb
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/flava/__init__.py
@@ -0,0 +1,31 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_flava import *
+ from .feature_extraction_flava import *
+ from .image_processing_flava import *
+ from .image_processing_pil_flava import *
+ from .modeling_flava import *
+ from .processing_flava import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/flava/configuration_flava.py b/third_party/transformers/src/transformers/models/flava/configuration_flava.py
new file mode 100644
index 0000000000000000000000000000000000000000..a787a8b305ecf38ea925f8b81c5957bb0815f38f
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/flava/configuration_flava.py
@@ -0,0 +1,427 @@
+# Copyright 2022 Meta Platforms authors and The HuggingFace Team. All rights reserved.
+#
+# 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.
+"""FLAVA model configurations"""
+
+from typing import Any
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="facebook/flava-full")
+@strict
+class FlavaImageConfig(PreTrainedConfig):
+ r"""
+ mask_token (`bool`, *optional*, defaults to `True`):
+ Whether to use a mask token or not. Used in MIM (Masked Image Modeling) loss for FLAVA.
+
+ Example:
+
+ ```python
+ >>> from transformers import FlavaImageConfig, FlavaImageModel
+
+ >>> # Initializing a FlavaImageModel with style configuration
+ >>> configuration = FlavaImageConfig()
+
+ >>> # Initializing a FlavaImageModel model (with random weights) from the style configuration
+ >>> model = FlavaImageModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "flava_image_model"
+ base_config_key = "image_config"
+
+ hidden_size: int = 768
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 12
+ intermediate_size: int = 3072
+ hidden_act: str = "gelu"
+ hidden_dropout_prob: float | int = 0.0
+ attention_probs_dropout_prob: float | int = 0.0
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-12
+ image_size: int | list[int] | tuple[int, int] = 224
+ patch_size: int | list[int] | tuple[int, int] = 16
+ num_channels: int = 3
+ qkv_bias: bool = True
+ mask_token: bool = True
+ vocab_size: int = 8192
+
+
+@auto_docstring(checkpoint="facebook/flava-full")
+@strict
+class FlavaTextConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import FlavaTextConfig, FlavaTextModel
+
+ >>> # Initializing a FlavaTextModel with style configuration
+ >>> configuration = FlavaTextConfig()
+
+ >>> # Initializing a FlavaTextModel model (with random weights) from the style configuration
+ >>> model = FlavaTextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "flava_text_model"
+ base_config_key = "text_config"
+
+ vocab_size: int = 30522
+ type_vocab_size: int = 2
+ max_position_embeddings: int = 512
+ hidden_size: int = 768
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 12
+ intermediate_size: int = 3072
+ hidden_act: str = "gelu"
+ hidden_dropout_prob: float | int = 0.0
+ attention_probs_dropout_prob: float | int = 0.0
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-12
+ pad_token_id: int | None = 0
+ qkv_bias: bool = True
+
+
+@auto_docstring(checkpoint="facebook/flava-full")
+@strict
+class FlavaMultimodalConfig(PreTrainedConfig):
+ r"""
+ use_cls_token (`bool`, *optional*, defaults to `True`):
+ Whether to use an extra CLS token for multimodal settings. Usually needed by the FLAVA model.
+
+ Example:
+
+ ```python
+ >>> from transformers import FlavaMultimodalConfig, FlavaMultimodalModel
+
+ >>> # Initializing a FlavaMultimodalModel with style configuration
+ >>> configuration = FlavaMultimodalConfig()
+
+ >>> # Initializing a FlavaMultimodalModel model (with random weights) from the style configuration
+ >>> model = FlavaMultimodalModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "flava_multimodal_model"
+ base_config_key = "multimodal_config"
+
+ hidden_size: int = 768
+ num_hidden_layers: int = 6
+ num_attention_heads: int = 12
+ intermediate_size: int = 3072
+ hidden_act: str = "gelu"
+ hidden_dropout_prob: float | int = 0.0
+ attention_probs_dropout_prob: float | int = 0.0
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-12
+ qkv_bias: bool = True
+ use_cls_token: bool = True
+
+
+@auto_docstring(checkpoint="facebook/flava-full")
+@strict
+class FlavaImageCodebookConfig(PreTrainedConfig):
+ r"""
+ num_groups (`int`, *optional*, defaults to 4):
+ Number of groups to be created. This parameter as of now doesn't affect the model and is used for some
+ internal calculation and estimations.
+ num_blocks_per_group (`int`, *optional*, defaults to 2):
+ Number of conv-based blocks per group.
+ freeze (`bool`, defaults to `True`):
+ Whether to freeze the weights of the model.
+
+ Example:
+
+ ```python
+ >>> from transformers import FlavaImageCodebookConfig, FlavaImageCodebook
+
+ >>> # Initializing a FlavaImageCodebook with style configuration
+ >>> configuration = FlavaImageCodebookConfig()
+
+ >>> # Initializing a FlavaImageCodebook model (with random weights) from the style configuration
+ >>> model = FlavaImageCodebook(configuration)
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+ """
+
+ num_groups: int = 4
+ input_channels: int = 3
+ num_blocks_per_group: int = 2
+ hidden_size: int = 256
+ vocab_size: int = 8192
+ freeze: bool = True
+ initializer_range: float = 0.02
+
+
+@auto_docstring(checkpoint="facebook/flava-full")
+@strict
+class FlavaConfig(PreTrainedConfig):
+ r"""
+ image_config (`dict`, *optional*):
+ Dictionary of configuration options used to initialize [`FlavaImageConfig`].
+ multimodal_config (`dict`, *optional*):
+ Dictionary of configuration options used to initialize [`FlavaMultimodalConfig`].
+ image_codebook_config (`dict`, *optional*):
+ Dictionary of configuration options used to initialize [`FlavaCodebookConfig`].
+ init_codebook (`bool`, *optional*, defaults to `True`):
+ Whether to initialize the codebook
+ logit_scale_init_value (`float`, *optional*, defaults to 2.6592):
+ The initial value of the *logit_scale* parameter. Default is used as per the original FLAVA/CLIP
+ implementation.
+ ce_ignore_index (`int`, *optional*, defaults to -100):
+ Cross entropy index to ignore.
+ mim_weight (`float`, *optional*, defaults to 1.0):
+ Weight to be assigned to MIM (Masked Image Modeling) unimodal loss
+ mlm_weight (`float`, *optional*, defaults to 1.0):
+ Weight to be assigned to MLM (Masked Language Modeling) unimodal loss
+ global_contrastive_weight (`float`, *optional*, defaults to 1.0):
+ Weight to be assigned to global contrastive cross-alignment loss.
+ itm_weight (`float`, *optional*, defaults to 1.0):
+ Weight to be assigned to image-text matching multimodal loss.
+ mmm_image_weight (`float`, *optional*, defaults to 1.0):
+ Weight to be assigned to MMM loss's image part.
+ mmm_text_weight (`float`, *optional*, defaults to 1.0):
+ Weight to be assigned to MMM loss's text part.
+ global_backprop_contrastive (`bool`, *optional*, defaults to `True`):
+ Whether to use global backpropgation through all workers in contrastive loss.
+ skip_unmasked_multimodal_encoder (`bool`, *optional*, defaults to `True`):
+ Whether to skip running unmasked multimodal encoder whose outputs are not used by FLAVA losses.
+ return_loss (`bool`, *optional*, defaults to `True`):
+ Whether to return loss or not
+
+ Example:
+
+ ```python
+ >>> from transformers import FlavaConfig, FlavaModel, FlavaForPreTraining
+
+ >>> # Initializing a FlavaConfig with style configuration
+ >>> configuration = FlavaConfig()
+
+ >>> # Initializing a FlavaModel and FlavaForPreTraining model (with random weights) from the style configuration
+ >>> model = FlavaModel(configuration)
+ >>> model_pre = FlavaForPreTraining(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ >>> configuration_pre = model_pre.config
+ ```
+ """
+
+ model_type = "flava"
+ sub_configs = {
+ "text_config": FlavaTextConfig,
+ "image_config": FlavaImageConfig,
+ "multimodal_config": FlavaMultimodalConfig,
+ "image_codebook_config": FlavaImageCodebookConfig,
+ }
+
+ image_config: dict[str, Any] | PreTrainedConfig | None = None
+ text_config: dict[str, Any] | PreTrainedConfig | None = None
+ multimodal_config: dict[str, Any] | PreTrainedConfig | None = None
+ image_codebook_config: dict[str, Any] | PreTrainedConfig | None = None
+ hidden_size: int = 768
+ layer_norm_eps: float = 1e-12
+ projection_dim: int = 768
+ init_codebook: bool = True
+ logit_scale_init_value: float = 2.6592
+ initializer_range: float = 0.02
+ ce_ignore_index: int = -100
+ mim_weight: float = 1.0
+ mlm_weight: float = 1.0
+ global_contrastive_weight: float = 1.0
+ itm_weight: float = 1.0
+ mmm_image_weight: float = 1.0
+ mmm_text_weight: float = 1.0
+ global_backprop_contrastive: bool = True
+ skip_unmasked_multimodal_encoder: bool = True
+ return_loss: bool = True
+ tie_word_embeddings: bool = True
+ initializer_factor: float = 1.0
+
+ def __post_init__(self, **kwargs):
+ if self.text_config is None:
+ text_config = {}
+ logger.info("`text_config` is `None`. Initializing the `FlavaTextConfig` with default values.")
+ elif isinstance(self.text_config, FlavaTextConfig):
+ text_config = self.text_config.to_dict()
+ else:
+ text_config = self.text_config
+
+ if self.image_config is None:
+ image_config = {}
+ logger.info("`image_config` is `None`. initializing the `FlavaImageConfig` with default values.")
+ elif isinstance(self.image_config, FlavaImageConfig):
+ image_config = self.image_config.to_dict()
+ else:
+ image_config = self.image_config
+
+ if self.multimodal_config is None:
+ multimodal_config = {}
+ logger.info("`multimodal_config` is `None`. Initializing the `FlavaMultimodalConfig` with default values.")
+ elif isinstance(self.multimodal_config, FlavaMultimodalConfig):
+ multimodal_config = self.multimodal_config.to_dict()
+ else:
+ multimodal_config = self.multimodal_config
+
+ if self.image_codebook_config is None:
+ image_codebook_config = {}
+ logger.info(
+ "`image_codebook_config` is `None`. initializing the `FlavaImageCodebookConfig` with default values."
+ )
+ elif isinstance(self.image_codebook_config, FlavaImageCodebookConfig):
+ image_codebook_config = self.image_codebook_config.to_dict()
+ else:
+ image_codebook_config = self.image_codebook_config
+
+ # If `_config_dict` exist, we use them for the backward compatibility.
+ text_config_dict = kwargs.pop("text_config_dict", None)
+ image_config_dict = kwargs.pop("image_config_dict", None)
+ multimodal_config_dict = kwargs.pop("multimodal_config_dict", None)
+ image_codebook_config_dict = kwargs.pop("image_codebook_config_dict", None)
+
+ # Instead of simply assigning `[text|vision]_config_dict` to `[text|vision]_config`, we use the values in
+ # `[text|vision]_config_dict` to update the values in `[text|vision]_config`. The values should be same in most
+ # cases, but we don't want to break anything regarding `_config_dict` that existed before commit `8827e1b2`.
+ if text_config_dict is not None:
+ # This is the complete result when using `text_config_dict`.
+ _text_config_dict = FlavaTextConfig(**text_config_dict).to_dict()
+
+ # Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
+ for key, value in _text_config_dict.items():
+ if key in text_config and value != text_config[key] and key != "transformers_version":
+ # If specified in `text_config_dict`
+ if key in text_config_dict:
+ message = (
+ f"`{key}` is found in both `text_config_dict` and `text_config` but with different values. "
+ f'The value `text_config_dict["{key}"]` will be used instead.'
+ )
+ # If inferred from default argument values (just to be super careful)
+ else:
+ message = (
+ f"`text_config_dict` is provided which will be used to initialize `FlavaTextConfig`. The "
+ f'value `text_config["{key}"]` will be overridden.'
+ )
+ logger.info(message)
+
+ # Update all values in `text_config` with the ones in `_text_config_dict`.
+ text_config.update(_text_config_dict)
+
+ if image_config_dict is not None:
+ # This is the complete result when using `image_config_dict`.
+ _image_config_dict = FlavaImageConfig(**image_config_dict).to_dict()
+ # convert keys to string instead of integer
+ if "id2label" in _image_config_dict:
+ _image_config_dict["id2label"] = {
+ str(key): value for key, value in _image_config_dict["id2label"].items()
+ }
+
+ # Give a warning if the values exist in both `_image_config_dict` and `image_config` but being different.
+ for key, value in _image_config_dict.items():
+ if key in image_config and value != image_config[key] and key != "transformers_version":
+ # If specified in `image_config_dict`
+ if key in image_config_dict:
+ message = (
+ f"`{key}` is found in both `image_config_dict` and `image_config` but with different "
+ f'values. The value `image_config_dict["{key}"]` will be used instead.'
+ )
+ # If inferred from default argument values (just to be super careful)
+ else:
+ message = (
+ f"`image_config_dict` is provided which will be used to initialize `FlavaImageConfig`. "
+ f'The value `image_config["{key}"]` will be overridden.'
+ )
+ logger.info(message)
+
+ # Update all values in `image_config` with the ones in `_image_config_dict`.
+ image_config.update(_image_config_dict)
+
+ if multimodal_config_dict is not None:
+ # This is the complete result when using `multimodal_config_dict`.
+ _multimodal_config_dict = FlavaMultimodalConfig(**multimodal_config_dict).to_dict()
+
+ # Give a warning if the values exist in both `_multimodal_config_dict` and `multimodal_config` but being
+ # different.
+ for key, value in _multimodal_config_dict.items():
+ if key in multimodal_config and value != multimodal_config[key] and key != "transformers_version":
+ # If specified in `multimodal_config_dict`
+ if key in multimodal_config_dict:
+ message = (
+ f"`{key}` is found in both `multimodal_config_dict` and `multimodal_config` but with "
+ f'different values. The value `multimodal_config_dict["{key}"]` will be used instead.'
+ )
+ # If inferred from default argument values (just to be super careful)
+ else:
+ message = (
+ f"`multimodal_config_dict` is provided which will be used to initialize "
+ f'`FlavaMultimodalConfig`. The value `multimodal_config["{key}"]` will be overridden.'
+ )
+ logger.info(message)
+
+ # Update all values in `multimodal_config` with the ones in `_multimodal_config_dict`.
+ multimodal_config.update(_multimodal_config_dict)
+
+ if image_codebook_config_dict is not None:
+ # This is the complete result when using `image_codebook_config_dict`.
+ _image_codebook_config_dict = FlavaImageCodebookConfig(**image_codebook_config_dict).to_dict()
+
+ # Give a warning if the values exist in both `_image_codebook_config_dict` and `image_codebook_config` but
+ # being different.
+ for key, value in _image_codebook_config_dict.items():
+ if (
+ key in image_codebook_config
+ and value != image_codebook_config[key]
+ and key != "transformers_version"
+ ):
+ # If specified in `image_codebook_config_dict`
+ if key in image_codebook_config_dict:
+ message = (
+ f"`{key}` is found in both `image_codebook_config_dict` and `image_codebook_config` but "
+ f'with different values. The value `image_codebook_config_dict["{key}"]` will be used '
+ "instead."
+ )
+ # If inferred from default argument values (just to be super careful)
+ else:
+ message = (
+ f"`image_codebook_config_dict` is provided which will be used to initialize "
+ f'`FlavaImageCodebookConfig`. The value `image_codebook_config["{key}"]` will be overridden.'
+ )
+ logger.info(message)
+
+ # Update all values in `image_codebook_config` with the ones in `_image_codebook_config_dict`.
+ image_codebook_config.update(_image_codebook_config_dict)
+
+ # Finally we can convert back our unified text/vision configs to `PretrainedConfig`
+ self.text_config = FlavaTextConfig(**text_config)
+ self.image_config = FlavaImageConfig(**image_config)
+ self.multimodal_config = FlavaMultimodalConfig(**multimodal_config)
+ self.image_codebook_config = FlavaImageCodebookConfig(**image_codebook_config)
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["FlavaConfig", "FlavaImageCodebookConfig", "FlavaImageConfig", "FlavaMultimodalConfig", "FlavaTextConfig"]
diff --git a/third_party/transformers/src/transformers/models/flava/convert_dalle_to_flava_codebook.py b/third_party/transformers/src/transformers/models/flava/convert_dalle_to_flava_codebook.py
new file mode 100644
index 0000000000000000000000000000000000000000..a86675074533eeaf8cab2a406919446a6d49d834
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/flava/convert_dalle_to_flava_codebook.py
@@ -0,0 +1,101 @@
+# Copyright 2022 Meta Platforms authors and The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+import argparse
+import os
+
+import torch
+
+from transformers import FlavaImageCodebook, FlavaImageCodebookConfig
+
+
+def rreplace(s, old, new, occurrence):
+ li = s.rsplit(old, occurrence)
+ return new.join(li)
+
+
+def count_parameters(state_dict):
+ # encoder.embeddings are double copied in original FLAVA
+ return sum(param.float().sum() if "encoder.embeddings" not in key else 0 for key, param in state_dict.items())
+
+
+def upgrade_state_dict(state_dict):
+ upgrade = {}
+
+ group_keys = ["group_1", "group_2", "group_3", "group_4"]
+ for key, value in state_dict.items():
+ for group_key in group_keys:
+ if group_key in key:
+ key = key.replace(f"{group_key}.", f"{group_key}.group.")
+
+ if "res_path" in key:
+ key = key.replace("res_path.", "res_path.path.")
+
+ if key.endswith(".w"):
+ key = rreplace(key, ".w", ".weight", 1)
+ if key.endswith(".b"):
+ key = rreplace(key, ".b", ".bias", 1)
+
+ upgrade[key] = value.float()
+
+ return upgrade
+
+
+@torch.no_grad()
+def convert_dalle_checkpoint(checkpoint_path, pytorch_dump_folder_path, config_path=None, save_checkpoint=True):
+ """
+ Copy/paste/tweak model's weights to transformers design.
+ """
+ from dall_e import Encoder
+
+ encoder = Encoder()
+ if os.path.exists(checkpoint_path):
+ ckpt = torch.load(checkpoint_path, weights_only=True)
+ else:
+ ckpt = torch.hub.load_state_dict_from_url(checkpoint_path)
+
+ if isinstance(ckpt, Encoder):
+ ckpt = ckpt.state_dict()
+ encoder.load_state_dict(ckpt)
+
+ if config_path is not None:
+ config = FlavaImageCodebookConfig.from_pretrained(config_path)
+ else:
+ config = FlavaImageCodebookConfig()
+
+ hf_model = FlavaImageCodebook(config).eval()
+ state_dict = encoder.state_dict()
+
+ hf_state_dict = upgrade_state_dict(state_dict)
+ hf_model.load_state_dict(hf_state_dict)
+ hf_state_dict = hf_model.state_dict()
+ hf_count = count_parameters(hf_state_dict)
+ state_dict_count = count_parameters(state_dict)
+
+ assert torch.allclose(hf_count, state_dict_count, atol=1e-3)
+
+ if save_checkpoint:
+ hf_model.save_pretrained(pytorch_dump_folder_path)
+ else:
+ return hf_state_dict
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
+ parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to flava checkpoint")
+ parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
+ args = parser.parse_args()
+
+ convert_dalle_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)
diff --git a/third_party/transformers/src/transformers/models/flava/convert_flava_original_pytorch_to_hf.py b/third_party/transformers/src/transformers/models/flava/convert_flava_original_pytorch_to_hf.py
new file mode 100644
index 0000000000000000000000000000000000000000..e621585a5750beb2a2dc5975b5b6c5473c8b272f
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/flava/convert_flava_original_pytorch_to_hf.py
@@ -0,0 +1,98 @@
+# Copyright 2022 Meta Platforms authors and The HuggingFace Team. All rights reserved.
+#
+# 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.
+
+import argparse
+import os
+
+import torch
+
+from transformers import FlavaConfig, FlavaForPreTraining
+from transformers.models.flava.convert_dalle_to_flava_codebook import convert_dalle_checkpoint
+
+
+def count_parameters(state_dict):
+ # encoder.embeddings are double copied in original FLAVA
+ return sum(param.float().sum() if "encoder.embeddings" not in key else 0 for key, param in state_dict.items())
+
+
+def upgrade_state_dict(state_dict, codebook_state_dict):
+ upgrade = {}
+
+ for key, value in state_dict.items():
+ if "text_encoder.embeddings" in key or "image_encoder.embeddings" in key:
+ continue
+
+ key = key.replace("heads.cmd.mim_head.cls.predictions", "mmm_image_head")
+ key = key.replace("heads.cmd.mlm_head.cls.predictions", "mmm_text_head")
+ key = key.replace("heads.cmd.itm_head.cls", "itm_head")
+ key = key.replace("heads.cmd.itm_head.pooler", "itm_head.pooler")
+ key = key.replace("heads.cmd.clip_head.logit_scale", "flava.logit_scale")
+ key = key.replace("heads.fairseq_mlm.cls.predictions", "mlm_head")
+ key = key.replace("heads.imagenet.mim_head.cls.predictions", "mim_head")
+ key = key.replace("mm_text_projection", "flava.text_to_mm_projection")
+ key = key.replace("mm_image_projection", "flava.image_to_mm_projection")
+ key = key.replace("image_encoder.module", "flava.image_model")
+ key = key.replace("text_encoder.module", "flava.text_model")
+ key = key.replace("mm_encoder.module.encoder.cls_token", "flava.multimodal_model.cls_token")
+ key = key.replace("mm_encoder.module", "flava.multimodal_model")
+ key = key.replace("text_projection", "flava.text_projection")
+ key = key.replace("image_projection", "flava.image_projection")
+
+ upgrade[key] = value.float()
+
+ for key, value in codebook_state_dict.items():
+ upgrade[f"image_codebook.{key}"] = value
+
+ return upgrade
+
+
+@torch.no_grad()
+def convert_flava_checkpoint(checkpoint_path, codebook_path, pytorch_dump_folder_path, config_path=None):
+ """
+ Copy/paste/tweak model's weights to transformers design.
+ """
+ if config_path is not None:
+ config = FlavaConfig.from_pretrained(config_path)
+ else:
+ config = FlavaConfig()
+
+ hf_model = FlavaForPreTraining(config).eval()
+
+ codebook_state_dict = convert_dalle_checkpoint(codebook_path, None, save_checkpoint=False)
+
+ if os.path.exists(checkpoint_path):
+ state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+ else:
+ state_dict = torch.hub.load_state_dict_from_url(checkpoint_path, map_location="cpu")
+
+ hf_state_dict = upgrade_state_dict(state_dict, codebook_state_dict)
+ hf_model.load_state_dict(hf_state_dict)
+ hf_state_dict = hf_model.state_dict()
+ hf_count = count_parameters(hf_state_dict)
+ state_dict_count = count_parameters(state_dict) + count_parameters(codebook_state_dict)
+
+ assert torch.allclose(hf_count, state_dict_count, atol=1e-3)
+
+ hf_model.save_pretrained(pytorch_dump_folder_path)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
+ parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to flava checkpoint")
+ parser.add_argument("--codebook_path", default=None, type=str, help="Path to flava codebook checkpoint")
+ parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
+ args = parser.parse_args()
+
+ convert_flava_checkpoint(args.checkpoint_path, args.codebook_path, args.pytorch_dump_folder_path, args.config_path)
diff --git a/third_party/transformers/src/transformers/models/flava/image_processing_flava.py b/third_party/transformers/src/transformers/models/flava/image_processing_flava.py
new file mode 100644
index 0000000000000000000000000000000000000000..bab4b1113e3ba2ff6a0f9867e0434077dc933516
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/flava/image_processing_flava.py
@@ -0,0 +1,475 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Image processor class for Flava."""
+
+import math
+import random
+from collections.abc import Iterable
+from functools import lru_cache
+from typing import Any
+
+import torch
+from torchvision.transforms.v2 import functional as tvF
+
+from ...image_processing_backends import TorchvisionBackend
+from ...image_processing_utils import BatchFeature, get_size_dict
+from ...image_transforms import group_images_by_shape, reorder_images
+from ...image_utils import (
+ OPENAI_CLIP_MEAN,
+ OPENAI_CLIP_STD,
+ ChannelDimension,
+ ImageInput,
+ PILImageResampling,
+ SizeDict,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import (
+ TensorType,
+ auto_docstring,
+)
+
+
+# These values are taken from CLIP
+FLAVA_IMAGE_MEAN = OPENAI_CLIP_MEAN
+FLAVA_IMAGE_STD = OPENAI_CLIP_STD
+FLAVA_CODEBOOK_MEAN = [0.0, 0.0, 0.0]
+FLAVA_CODEBOOK_STD = [1.0, 1.0, 1.0]
+LOGIT_LAPLACE_EPS: float = 0.1
+
+
+class FlavaImageProcessorKwargs(ImagesKwargs, total=False):
+ """
+ return_image_mask (`bool`, *optional*, defaults to `False`):
+ Whether to return the image mask. Can be overridden by the `return_image_mask` parameter in `preprocess`.
+ input_size_patches (`int`, *optional*, defaults to `14`):
+ Number of patches in the image in height and width direction. 14x14 = 196 total patches. Can be overridden
+ by the `input_size_patches` parameter in `preprocess`.
+ total_mask_patches (`int`, *optional*, defaults to `75`):
+ Total number of patches that should be masked. Can be overridden by the `total_mask_patches` parameter in
+ `preprocess`.
+ mask_group_min_patches (`int`, *optional*, defaults to `16`):
+ Minimum number of patches that should be masked. Can be overridden by the `mask_group_min_patches`
+ parameter in `preprocess`.
+ mask_group_max_patches (`int`, *optional*):
+ Maximum number of patches that should be masked. Can be overridden by the `mask_group_max_patches`
+ parameter in `preprocess`.
+ mask_group_min_aspect_ratio (`float`, *optional*, defaults to `0.3`):
+ Minimum aspect ratio of the mask window. Can be overridden by the `mask_group_min_aspect_ratio` parameter
+ in `preprocess`.
+ mask_group_max_aspect_ratio (`float`, *optional*):
+ Maximum aspect ratio of the mask window. Can be overridden by the `mask_group_max_aspect_ratio` parameter
+ in `preprocess`.
+ return_codebook_pixels (`bool`, *optional*, defaults to `False`):
+ Whether to return the codebook pixel values.
+ codebook_do_resize (`bool`, *optional*, defaults to `True`):
+ Whether to resize the input for codebook to a certain. Can be overridden by the `codebook_do_resize`
+ parameter in `preprocess`. `codebook_size`.
+ codebook_size (`dict[str, int]`, *optional*, defaults to `{"height": 224, "width": 224}`):
+ Resize the input for codebook to the given size. Can be overridden by the `codebook_size` parameter in
+ `preprocess`.
+ codebook_resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.LANCZOS` for PIL backend,
+ `PILImageResampling.BICUBIC` for torchvision backend):
+ Resampling filter to use if resizing the codebook image. LANCZOS is not supported for torch Tensors;
+ BICUBIC is used as the closest alternative for the torchvision backend. Can be overridden by the
+ `codebook_resample` parameter in `preprocess`.
+ codebook_do_center_crop (`bool`, *optional*, defaults to `True`):
+ Whether to crop the input for codebook at the center. If the input size is smaller than
+ `codebook_crop_size` along any edge, the image is padded with 0's and then center cropped. Can be
+ overridden by the `codebook_do_center_crop` parameter in `preprocess`.
+ codebook_crop_size (`dict[str, int]`, *optional*, defaults to `{"height": 224, "width": 224}`):
+ Desired output size for codebook input when applying center-cropping. Can be overridden by the
+ `codebook_crop_size` parameter in `preprocess`.
+ codebook_do_rescale (`bool`, *optional*, defaults to `True`):
+ Whether to rescale the input for codebook by the specified scale `codebook_rescale_factor`. Can be
+ overridden by the `codebook_do_rescale` parameter in `preprocess`.
+ codebook_rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
+ Defines the scale factor to use if rescaling the codebook image. Can be overridden by the
+ `codebook_rescale_factor` parameter in `preprocess`.
+ codebook_do_map_pixels (`bool`, *optional*, defaults to `True`):
+ Whether to map the pixel values of the codebook input to (1 - 2e)x + e. Can be overridden by the
+ `codebook_do_map_pixels` parameter in `preprocess`.
+ codebook_do_normalize (`bool`, *optional*, defaults to `True`):
+ Whether or not to normalize the input for codebook with `codebook_image_mean` and `codebook_image_std`. Can
+ be overridden by the `codebook_do_normalize` parameter in `preprocess`.
+ codebook_image_mean (`Optional[Union[float, Iterable[float]]]`, *optional*, defaults to `[0, 0, 0]`):
+ The sequence of means for each channel, to be used when normalizing images for codebook. Can be overridden
+ by the `codebook_image_mean` parameter in `preprocess`.
+ codebook_image_std (`Optional[Union[float, Iterable[float]]]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):
+ The sequence of standard deviations for each channel, to be used when normalizing images for codebook. Can
+ be overridden by the `codebook_image_std` parameter in `preprocess`.
+ """
+
+ # Mask related params
+ return_image_mask: bool
+ input_size_patches: int
+ total_mask_patches: int
+ mask_group_min_patches: int
+ mask_group_max_patches: int
+ mask_group_min_aspect_ratio: float
+ mask_group_max_aspect_ratio: float
+ # Codebook related params
+ return_codebook_pixels: bool
+ codebook_do_resize: bool
+ codebook_size: dict[str, int]
+ codebook_resample: int
+ codebook_do_center_crop: bool
+ codebook_crop_size: dict[str, int]
+ codebook_do_rescale: bool
+ codebook_rescale_factor: int | float
+ codebook_do_map_pixels: bool
+ codebook_do_normalize: bool
+ codebook_image_mean: float | Iterable[float]
+ codebook_image_std: float | Iterable[float]
+
+
+# Inspired from https://github.com/microsoft/unilm/blob/master/beit/masking_generator.py
+class FlavaMaskingGenerator:
+ def __init__(
+ self,
+ input_size: int | tuple[int, int] = 14,
+ total_mask_patches: int = 75,
+ mask_group_max_patches: int | None = None,
+ mask_group_min_patches: int = 16,
+ mask_group_min_aspect_ratio: float | None = 0.3,
+ mask_group_max_aspect_ratio: float | None = None,
+ ):
+ if not isinstance(input_size, tuple):
+ input_size = (input_size,) * 2
+ self.height, self.width = input_size
+
+ self.num_patches = self.height * self.width
+ self.total_mask_patches = total_mask_patches
+
+ self.mask_group_min_patches = mask_group_min_patches
+ self.mask_group_max_patches = total_mask_patches if mask_group_max_patches is None else mask_group_max_patches
+
+ mask_group_max_aspect_ratio = mask_group_max_aspect_ratio or 1 / mask_group_min_aspect_ratio
+ self.log_aspect_ratio = (math.log(mask_group_min_aspect_ratio), math.log(mask_group_max_aspect_ratio))
+
+ def __repr__(self):
+ repr_str = "MaskingGenerator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" % (
+ self.height,
+ self.width,
+ self.mask_group_min_patches,
+ self.mask_group_max_patches,
+ self.total_mask_patches,
+ self.log_aspect_ratio[0],
+ self.log_aspect_ratio[1],
+ )
+ return repr_str
+
+ def get_shape(self):
+ return self.height, self.width
+
+ def _mask(self, mask, max_mask_patches):
+ delta = 0
+ for _attempt in range(10):
+ target_area = random.uniform(self.mask_group_min_patches, max_mask_patches)
+ aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
+ height = int(round(math.sqrt(target_area * aspect_ratio)))
+ width = int(round(math.sqrt(target_area / aspect_ratio)))
+ if width < self.width and height < self.height:
+ top = random.randint(0, self.height - height)
+ left = random.randint(0, self.width - width)
+
+ num_masked = mask[top : top + height, left : left + width].sum()
+ # Overlap
+ if 0 < height * width - num_masked <= max_mask_patches:
+ zeros_pos = mask[top : top + height, left : left + width] == 0
+ mask[top : top + height, left : left + width][zeros_pos] = 1
+ delta += zeros_pos.sum()
+
+ if delta > 0:
+ break
+ return delta
+
+ def __call__(self):
+ mask = torch.zeros(self.get_shape(), dtype=torch.int)
+ mask_count = 0
+ while mask_count < self.total_mask_patches:
+ max_mask_patches = self.total_mask_patches - mask_count
+ max_mask_patches = min(max_mask_patches, self.mask_group_max_patches)
+
+ delta = self._mask(mask, max_mask_patches)
+ if delta == 0:
+ break
+ else:
+ mask_count += delta
+
+ return mask
+
+
+@auto_docstring
+class FlavaImageProcessor(TorchvisionBackend):
+ valid_kwargs = FlavaImageProcessorKwargs
+ resample = PILImageResampling.BICUBIC
+ image_mean = FLAVA_IMAGE_MEAN
+ image_std = FLAVA_IMAGE_STD
+ size = {"height": 224, "width": 224}
+ crop_size = {"height": 224, "width": 224}
+ do_resize = True
+ do_center_crop = True
+ do_rescale = True
+ do_normalize = True
+
+ # Mask related params
+ return_image_mask = False
+ input_size_patches = 14
+ total_mask_patches = 75
+ mask_group_min_patches = 16
+ mask_group_max_patches = None
+ mask_group_min_aspect_ratio = 0.3
+ mask_group_max_aspect_ratio = None
+ # Codebook related params
+ return_codebook_pixels = False
+ codebook_do_resize = True
+ codebook_size = {"height": 112, "width": 112}
+ # LANCZOS resampling is not supported for torch Tensors; BICUBIC is the closest alternative.
+ # Note: the PIL backend defaults to LANCZOS for this parameter.
+ codebook_resample = PILImageResampling.BICUBIC
+ codebook_do_center_crop = True
+ codebook_crop_size = {"height": 112, "width": 112}
+ codebook_do_rescale = True
+ codebook_rescale_factor = 1 / 255
+ codebook_do_map_pixels = True
+ codebook_do_normalize = True
+ codebook_image_mean = FLAVA_CODEBOOK_MEAN
+ codebook_image_std = FLAVA_CODEBOOK_STD
+
+ def __init__(self, **kwargs: Unpack[FlavaImageProcessorKwargs]):
+ super().__init__(**kwargs)
+
+ @auto_docstring
+ def preprocess(self, images: ImageInput, **kwargs: Unpack[FlavaImageProcessorKwargs]) -> BatchFeature:
+ return super().preprocess(images, **kwargs)
+
+ @classmethod
+ def from_dict(cls, image_processor_dict: dict[str, Any], **kwargs):
+ """
+ Overrides the `from_dict` method from the base class to make sure parameters are updated if image processor is
+ created using from_dict and kwargs e.g. `FlavaImageProcessor.from_pretrained(checkpoint, codebook_size=600)`
+ """
+ image_processor_dict = image_processor_dict.copy()
+ if "codebook_size" in kwargs:
+ image_processor_dict["codebook_size"] = kwargs.pop("codebook_size")
+ if "codebook_crop_size" in kwargs:
+ image_processor_dict["codebook_crop_size"] = kwargs.pop("codebook_crop_size")
+ return super().from_dict(image_processor_dict, **kwargs)
+
+ @lru_cache
+ def masking_generator(
+ self,
+ input_size_patches,
+ total_mask_patches,
+ mask_group_min_patches,
+ mask_group_max_patches,
+ mask_group_min_aspect_ratio,
+ mask_group_max_aspect_ratio,
+ ) -> FlavaMaskingGenerator:
+ return FlavaMaskingGenerator(
+ input_size=input_size_patches,
+ total_mask_patches=total_mask_patches,
+ mask_group_min_patches=mask_group_min_patches,
+ mask_group_max_patches=mask_group_max_patches,
+ mask_group_min_aspect_ratio=mask_group_min_aspect_ratio,
+ mask_group_max_aspect_ratio=mask_group_max_aspect_ratio,
+ )
+
+ def map_pixels(self, image: "torch.Tensor") -> "torch.Tensor":
+ return (1 - 2 * LOGIT_LAPLACE_EPS) * image + LOGIT_LAPLACE_EPS
+
+ def _standardize_kwargs(
+ self,
+ size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
+ crop_size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
+ default_to_square: bool | None = None,
+ image_mean: float | list[float] | None = None,
+ image_std: float | list[float] | None = None,
+ codebook_size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
+ codebook_crop_size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
+ codebook_image_mean: float | list[float] | None = None,
+ codebook_image_std: float | list[float] | None = None,
+ codebook_resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None,
+ data_format: ChannelDimension | None = None,
+ **kwargs,
+ ) -> dict:
+ """
+ Update kwargs that need further processing before being validated
+ Can be overridden by subclasses to customize the processing of kwargs.
+ """
+ kwargs = super()._standardize_kwargs(
+ size=size,
+ crop_size=crop_size,
+ default_to_square=default_to_square,
+ image_mean=image_mean,
+ image_std=image_std,
+ data_format=data_format,
+ **kwargs,
+ )
+ if codebook_size is not None and not isinstance(codebook_size, SizeDict):
+ codebook_size = SizeDict(**get_size_dict(size=codebook_size, default_to_square=default_to_square))
+ if codebook_crop_size is not None and not isinstance(codebook_crop_size, SizeDict):
+ codebook_crop_size = SizeDict(**get_size_dict(codebook_crop_size, param_name="codebook_crop_size"))
+ if isinstance(codebook_image_mean, list):
+ codebook_image_mean = tuple(codebook_image_mean)
+ if isinstance(codebook_image_std, list):
+ codebook_image_std = tuple(codebook_image_std)
+
+ kwargs["codebook_size"] = codebook_size
+ kwargs["codebook_crop_size"] = codebook_crop_size
+ kwargs["codebook_image_mean"] = codebook_image_mean
+ kwargs["codebook_image_std"] = codebook_image_std
+ # Store codebook_resample as-is - resize() will handle conversion
+ kwargs["codebook_resample"] = codebook_resample
+
+ return kwargs
+
+ def _preprocess_image(
+ self,
+ images: list["torch.Tensor"],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+ do_center_crop: bool,
+ crop_size: SizeDict,
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ do_map_pixels: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ disable_grouping: bool | None,
+ return_tensors: str | TensorType | None,
+ ) -> "torch.Tensor":
+ # Group images by size for batched resizing
+ grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+ resized_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ if do_resize:
+ stacked_images = self.resize(image=stacked_images, size=size, resample=resample)
+ resized_images_grouped[shape] = stacked_images
+ resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+ # Group images by size for further processing
+ # Needed in case do_resize is False, or resize returns images with different sizes
+ grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
+ processed_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ if do_center_crop:
+ stacked_images = self.center_crop(stacked_images, crop_size)
+ # Fused rescale and normalize
+ stacked_images = self.rescale_and_normalize(
+ stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
+ )
+ if do_map_pixels:
+ stacked_images = self.map_pixels(image=stacked_images)
+ processed_images_grouped[shape] = stacked_images
+
+ processed_images = reorder_images(processed_images_grouped, grouped_images_index)
+
+ return processed_images
+
+ def _preprocess(
+ self,
+ images: list["torch.Tensor"],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+ do_center_crop: bool,
+ crop_size: SizeDict,
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ # Mask related params
+ return_image_mask: bool | None,
+ input_size_patches: int | None,
+ total_mask_patches: int | None,
+ mask_group_min_patches: int | None,
+ mask_group_max_patches: int | None,
+ mask_group_min_aspect_ratio: float | None,
+ mask_group_max_aspect_ratio: float | None,
+ # Codebook related params
+ return_codebook_pixels: bool | None,
+ codebook_do_resize: bool | None,
+ codebook_size: SizeDict | None,
+ codebook_resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+ codebook_do_center_crop: bool | None,
+ codebook_crop_size: SizeDict | None,
+ codebook_do_rescale: bool | None,
+ codebook_rescale_factor: float | None,
+ codebook_do_map_pixels: bool | None,
+ codebook_do_normalize: bool | None,
+ codebook_image_mean: float | list[float] | None,
+ codebook_image_std: float | list[float] | None,
+ disable_grouping: bool | None,
+ return_tensors: str | TensorType | None,
+ **kwargs,
+ ) -> BatchFeature:
+ processed_images = self._preprocess_image(
+ images=images,
+ do_resize=do_resize,
+ size=size,
+ resample=resample,
+ do_center_crop=do_center_crop,
+ crop_size=crop_size,
+ do_rescale=do_rescale,
+ rescale_factor=rescale_factor,
+ do_normalize=do_normalize,
+ do_map_pixels=False,
+ image_mean=image_mean,
+ image_std=image_std,
+ disable_grouping=disable_grouping,
+ return_tensors=return_tensors,
+ )
+ data = {
+ "pixel_values": processed_images,
+ }
+
+ if return_codebook_pixels:
+ codebook_processed_images = self._preprocess_image(
+ images=images,
+ do_resize=codebook_do_resize,
+ size=codebook_size,
+ resample=codebook_resample, # resize() handles conversion from PILImageResampling to InterpolationMode
+ do_center_crop=codebook_do_center_crop,
+ crop_size=codebook_crop_size,
+ do_rescale=codebook_do_rescale,
+ rescale_factor=codebook_rescale_factor,
+ do_normalize=codebook_do_normalize,
+ do_map_pixels=codebook_do_map_pixels,
+ image_mean=codebook_image_mean,
+ image_std=codebook_image_std,
+ disable_grouping=disable_grouping,
+ return_tensors=return_tensors,
+ )
+ data["codebook_pixel_values"] = codebook_processed_images
+
+ if return_image_mask:
+ mask_generator = self.masking_generator(
+ input_size_patches=input_size_patches,
+ total_mask_patches=total_mask_patches,
+ mask_group_min_patches=mask_group_min_patches,
+ mask_group_max_patches=mask_group_max_patches,
+ mask_group_min_aspect_ratio=mask_group_min_aspect_ratio,
+ mask_group_max_aspect_ratio=mask_group_max_aspect_ratio,
+ )
+ masks = [mask_generator() for _ in range(len(images))]
+ data["bool_masked_pos"] = masks
+
+ return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["FlavaImageProcessor"]
diff --git a/third_party/transformers/src/transformers/models/flava/image_processing_pil_flava.py b/third_party/transformers/src/transformers/models/flava/image_processing_pil_flava.py
new file mode 100644
index 0000000000000000000000000000000000000000..ef2228d6ddf18af570d0e2236ed37a0ee852544a
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/flava/image_processing_pil_flava.py
@@ -0,0 +1,473 @@
+# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Image processor class for Flava."""
+
+import math
+import random
+from collections.abc import Iterable
+from functools import lru_cache
+from typing import Any
+
+import numpy as np
+
+from ...image_processing_backends import PilBackend
+from ...image_processing_utils import BatchFeature, get_size_dict
+from ...image_utils import (
+ OPENAI_CLIP_MEAN,
+ OPENAI_CLIP_STD,
+ ChannelDimension,
+ ImageInput,
+ PILImageResampling,
+ SizeDict,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import (
+ TensorType,
+ auto_docstring,
+ is_torch_available,
+)
+
+
+if is_torch_available():
+ import torch
+
+
+# Adapted from transformers.models.flava.image_processing_flava.FLAVA_CODEBOOK_MEAN
+FLAVA_CODEBOOK_MEAN = [0.0, 0.0, 0.0]
+
+# Adapted from transformers.models.flava.image_processing_flava.FLAVA_CODEBOOK_STD
+FLAVA_CODEBOOK_STD = [1.0, 1.0, 1.0]
+
+# Adapted from transformers.models.flava.image_processing_flava.FLAVA_IMAGE_MEAN
+# These values are taken from CLIP
+FLAVA_IMAGE_MEAN = OPENAI_CLIP_MEAN
+
+# Adapted from transformers.models.flava.image_processing_flava.FLAVA_IMAGE_STD
+FLAVA_IMAGE_STD = OPENAI_CLIP_STD
+
+# Adapted from transformers.models.flava.image_processing_flava.LOGIT_LAPLACE_EPS
+LOGIT_LAPLACE_EPS: float = 0.1
+
+
+# Adapted from transformers.models.flava.image_processing_flava.FlavaImageProcessorKwargs
+class FlavaImageProcessorKwargs(ImagesKwargs, total=False):
+ """
+ return_image_mask (`bool`, *optional*, defaults to `False`):
+ Whether to return the image mask. Can be overridden by the `return_image_mask` parameter in `preprocess`.
+ input_size_patches (`int`, *optional*, defaults to `14`):
+ Number of patches in the image in height and width direction. 14x14 = 196 total patches. Can be overridden
+ by the `input_size_patches` parameter in `preprocess`.
+ total_mask_patches (`int`, *optional*, defaults to `75`):
+ Total number of patches that should be masked. Can be overridden by the `total_mask_patches` parameter in
+ `preprocess`.
+ mask_group_min_patches (`int`, *optional*, defaults to `16`):
+ Minimum number of patches that should be masked. Can be overridden by the `mask_group_min_patches`
+ parameter in `preprocess`.
+ mask_group_max_patches (`int`, *optional*):
+ Maximum number of patches that should be masked. Can be overridden by the `mask_group_max_patches`
+ parameter in `preprocess`.
+ mask_group_min_aspect_ratio (`float`, *optional*, defaults to `0.3`):
+ Minimum aspect ratio of the mask window. Can be overridden by the `mask_group_min_aspect_ratio` parameter
+ in `preprocess`.
+ mask_group_max_aspect_ratio (`float`, *optional*):
+ Maximum aspect ratio of the mask window. Can be overridden by the `mask_group_max_aspect_ratio` parameter
+ in `preprocess`.
+ return_codebook_pixels (`bool`, *optional*, defaults to `False`):
+ Whether to return the codebook pixel values.
+ codebook_do_resize (`bool`, *optional*, defaults to `True`):
+ Whether to resize the input for codebook to a certain. Can be overridden by the `codebook_do_resize`
+ parameter in `preprocess`. `codebook_size`.
+ codebook_size (`dict[str, int]`, *optional*, defaults to `{"height": 224, "width": 224}`):
+ Resize the input for codebook to the given size. Can be overridden by the `codebook_size` parameter in
+ `preprocess`.
+ codebook_resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.LANCZOS` for PIL backend,
+ `PILImageResampling.BICUBIC` for torchvision backend):
+ Resampling filter to use if resizing the codebook image. LANCZOS is not supported for torch Tensors;
+ BICUBIC is used as the closest alternative for the torchvision backend. Can be overridden by the
+ `codebook_resample` parameter in `preprocess`.
+ codebook_do_center_crop (`bool`, *optional*, defaults to `True`):
+ Whether to crop the input for codebook at the center. If the input size is smaller than
+ `codebook_crop_size` along any edge, the image is padded with 0's and then center cropped. Can be
+ overridden by the `codebook_do_center_crop` parameter in `preprocess`.
+ codebook_crop_size (`dict[str, int]`, *optional*, defaults to `{"height": 224, "width": 224}`):
+ Desired output size for codebook input when applying center-cropping. Can be overridden by the
+ `codebook_crop_size` parameter in `preprocess`.
+ codebook_do_rescale (`bool`, *optional*, defaults to `True`):
+ Whether to rescale the input for codebook by the specified scale `codebook_rescale_factor`. Can be
+ overridden by the `codebook_do_rescale` parameter in `preprocess`.
+ codebook_rescale_factor (`int` or `float`, *optional*, defaults to `1/255`):
+ Defines the scale factor to use if rescaling the codebook image. Can be overridden by the
+ `codebook_rescale_factor` parameter in `preprocess`.
+ codebook_do_map_pixels (`bool`, *optional*, defaults to `True`):
+ Whether to map the pixel values of the codebook input to (1 - 2e)x + e. Can be overridden by the
+ `codebook_do_map_pixels` parameter in `preprocess`.
+ codebook_do_normalize (`bool`, *optional*, defaults to `True`):
+ Whether or not to normalize the input for codebook with `codebook_image_mean` and `codebook_image_std`. Can
+ be overridden by the `codebook_do_normalize` parameter in `preprocess`.
+ codebook_image_mean (`Optional[Union[float, Iterable[float]]]`, *optional*, defaults to `[0, 0, 0]`):
+ The sequence of means for each channel, to be used when normalizing images for codebook. Can be overridden
+ by the `codebook_image_mean` parameter in `preprocess`.
+ codebook_image_std (`Optional[Union[float, Iterable[float]]]`, *optional*, defaults to `[0.5, 0.5, 0.5]`):
+ The sequence of standard deviations for each channel, to be used when normalizing images for codebook. Can
+ be overridden by the `codebook_image_std` parameter in `preprocess`.
+ """
+
+ # Mask related params
+ return_image_mask: bool
+ input_size_patches: int
+ total_mask_patches: int
+ mask_group_min_patches: int
+ mask_group_max_patches: int
+ mask_group_min_aspect_ratio: float
+ mask_group_max_aspect_ratio: float
+ # Codebook related params
+ return_codebook_pixels: bool
+ codebook_do_resize: bool
+ codebook_size: dict[str, int]
+ codebook_resample: int
+ codebook_do_center_crop: bool
+ codebook_crop_size: dict[str, int]
+ codebook_do_rescale: bool
+ codebook_rescale_factor: int | float
+ codebook_do_map_pixels: bool
+ codebook_do_normalize: bool
+ codebook_image_mean: float | Iterable[float]
+ codebook_image_std: float | Iterable[float]
+
+
+# Adapted from transformers.models.flava.image_processing_flava.FlavaMaskingGenerator
+# Inspired from https://github.com/microsoft/unilm/blob/master/beit/masking_generator.py
+class FlavaMaskingGenerator:
+ def __init__(
+ self,
+ input_size: int | tuple[int, int] = 14,
+ total_mask_patches: int = 75,
+ mask_group_max_patches: int | None = None,
+ mask_group_min_patches: int = 16,
+ mask_group_min_aspect_ratio: float | None = 0.3,
+ mask_group_max_aspect_ratio: float | None = None,
+ ):
+ if not isinstance(input_size, tuple):
+ input_size = (input_size,) * 2
+ self.height, self.width = input_size
+
+ self.num_patches = self.height * self.width
+ self.total_mask_patches = total_mask_patches
+
+ self.mask_group_min_patches = mask_group_min_patches
+ self.mask_group_max_patches = total_mask_patches if mask_group_max_patches is None else mask_group_max_patches
+
+ mask_group_max_aspect_ratio = mask_group_max_aspect_ratio or 1 / mask_group_min_aspect_ratio
+ self.log_aspect_ratio = (math.log(mask_group_min_aspect_ratio), math.log(mask_group_max_aspect_ratio))
+
+ def __repr__(self):
+ repr_str = "MaskingGenerator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" % (
+ self.height,
+ self.width,
+ self.mask_group_min_patches,
+ self.mask_group_max_patches,
+ self.total_mask_patches,
+ self.log_aspect_ratio[0],
+ self.log_aspect_ratio[1],
+ )
+ return repr_str
+
+ def get_shape(self):
+ return self.height, self.width
+
+ def _mask(self, mask, max_mask_patches):
+ delta = 0
+ for _attempt in range(10):
+ target_area = random.uniform(self.mask_group_min_patches, max_mask_patches)
+ aspect_ratio = math.exp(random.uniform(*self.log_aspect_ratio))
+ height = int(round(math.sqrt(target_area * aspect_ratio)))
+ width = int(round(math.sqrt(target_area / aspect_ratio)))
+ if width < self.width and height < self.height:
+ top = random.randint(0, self.height - height)
+ left = random.randint(0, self.width - width)
+
+ num_masked = mask[top : top + height, left : left + width].sum()
+ # Overlap
+ if 0 < height * width - num_masked <= max_mask_patches:
+ zeros_pos = mask[top : top + height, left : left + width] == 0
+ mask[top : top + height, left : left + width][zeros_pos] = 1
+ delta += zeros_pos.sum()
+
+ if delta > 0:
+ break
+ return delta
+
+ def __call__(self):
+ mask = torch.zeros(self.get_shape(), dtype=torch.int)
+ mask_count = 0
+ while mask_count < self.total_mask_patches:
+ max_mask_patches = self.total_mask_patches - mask_count
+ max_mask_patches = min(max_mask_patches, self.mask_group_max_patches)
+
+ delta = self._mask(mask, max_mask_patches)
+ if delta == 0:
+ break
+ else:
+ mask_count += delta
+
+ return mask
+
+
+@auto_docstring
+class FlavaImageProcessorPil(PilBackend):
+ valid_kwargs = FlavaImageProcessorKwargs
+ resample = PILImageResampling.BICUBIC
+ image_mean = FLAVA_IMAGE_MEAN
+ image_std = FLAVA_IMAGE_STD
+ size = {"height": 224, "width": 224}
+ crop_size = {"height": 224, "width": 224}
+ do_resize = True
+ do_center_crop = True
+ do_rescale = True
+ do_normalize = True
+
+ # Mask related params
+ return_image_mask = False
+ input_size_patches = 14
+ total_mask_patches = 75
+ mask_group_min_patches = 16
+ mask_group_max_patches = None
+ mask_group_min_aspect_ratio = 0.3
+ mask_group_max_aspect_ratio = None
+ # Codebook related params
+ return_codebook_pixels = False
+ codebook_do_resize = True
+ codebook_size = {"height": 112, "width": 112}
+ codebook_resample = PILImageResampling.LANCZOS
+ codebook_do_center_crop = True
+ codebook_crop_size = {"height": 112, "width": 112}
+ codebook_do_rescale = True
+ codebook_rescale_factor = 1 / 255
+ codebook_do_map_pixels = True
+ codebook_do_normalize = True
+ codebook_image_mean = FLAVA_CODEBOOK_MEAN
+ codebook_image_std = FLAVA_CODEBOOK_STD
+
+ def __init__(self, **kwargs: Unpack[FlavaImageProcessorKwargs]):
+ super().__init__(**kwargs)
+
+ @auto_docstring
+ def preprocess(self, images: ImageInput, **kwargs: Unpack[FlavaImageProcessorKwargs]) -> BatchFeature:
+ return super().preprocess(images, **kwargs)
+
+ @classmethod
+ def from_dict(cls, image_processor_dict: dict[str, Any], **kwargs):
+ """
+ Overrides the `from_dict` method from the base class to make sure parameters are updated if image processor is
+ created using from_dict and kwargs e.g. `FlavaImageProcessor.from_pretrained(checkpoint, codebook_size=600)`
+ """
+ image_processor_dict = image_processor_dict.copy()
+ if "codebook_size" in kwargs:
+ image_processor_dict["codebook_size"] = kwargs.pop("codebook_size")
+ if "codebook_crop_size" in kwargs:
+ image_processor_dict["codebook_crop_size"] = kwargs.pop("codebook_crop_size")
+ return super().from_dict(image_processor_dict, **kwargs)
+
+ @lru_cache
+ def masking_generator(
+ self,
+ input_size_patches,
+ total_mask_patches,
+ mask_group_min_patches,
+ mask_group_max_patches,
+ mask_group_min_aspect_ratio,
+ mask_group_max_aspect_ratio,
+ ) -> FlavaMaskingGenerator:
+ # Import torch here since FlavaMaskingGenerator uses torch
+ return FlavaMaskingGenerator(
+ input_size=input_size_patches,
+ total_mask_patches=total_mask_patches,
+ mask_group_min_patches=mask_group_min_patches,
+ mask_group_max_patches=mask_group_max_patches,
+ mask_group_min_aspect_ratio=mask_group_min_aspect_ratio,
+ mask_group_max_aspect_ratio=mask_group_max_aspect_ratio,
+ )
+
+ def map_pixels(self, image: np.ndarray) -> np.ndarray:
+ return (1 - 2 * LOGIT_LAPLACE_EPS) * image + LOGIT_LAPLACE_EPS
+
+ def _standardize_kwargs(
+ self,
+ size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
+ crop_size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
+ default_to_square: bool | None = None,
+ image_mean: float | list[float] | None = None,
+ image_std: float | list[float] | None = None,
+ codebook_size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
+ codebook_crop_size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
+ codebook_image_mean: float | list[float] | None = None,
+ codebook_image_std: float | list[float] | None = None,
+ codebook_resample: "PILImageResampling | None" = None,
+ data_format: ChannelDimension | None = None,
+ **kwargs,
+ ) -> dict:
+ """
+ Update kwargs that need further processing before being validated
+ Can be overridden by subclasses to customize the processing of kwargs.
+ """
+ kwargs = super()._standardize_kwargs(
+ size=size,
+ crop_size=crop_size,
+ default_to_square=default_to_square,
+ image_mean=image_mean,
+ image_std=image_std,
+ data_format=data_format,
+ **kwargs,
+ )
+ if codebook_size is not None and not isinstance(codebook_size, SizeDict):
+ codebook_size = SizeDict(**get_size_dict(size=codebook_size, default_to_square=default_to_square))
+ if codebook_crop_size is not None and not isinstance(codebook_crop_size, SizeDict):
+ codebook_crop_size = SizeDict(**get_size_dict(codebook_crop_size, param_name="codebook_crop_size"))
+ if isinstance(codebook_image_mean, list):
+ codebook_image_mean = tuple(codebook_image_mean)
+ if isinstance(codebook_image_std, list):
+ codebook_image_std = tuple(codebook_image_std)
+
+ kwargs["codebook_size"] = codebook_size
+ kwargs["codebook_crop_size"] = codebook_crop_size
+ kwargs["codebook_image_mean"] = codebook_image_mean
+ kwargs["codebook_image_std"] = codebook_image_std
+ # Store codebook_resample as-is - resize() will handle conversion
+ kwargs["codebook_resample"] = codebook_resample
+
+ return kwargs
+
+ def _preprocess_image(
+ self,
+ images: list[np.ndarray],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | None",
+ do_center_crop: bool,
+ crop_size: SizeDict,
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ do_map_pixels: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ return_tensors: str | TensorType | None,
+ ) -> list[np.ndarray]:
+ # Process images one by one (no batching in PIL backend)
+ processed_images = []
+ for image in images:
+ if do_resize:
+ image = self.resize(image=image, size=size, resample=resample)
+ if do_center_crop:
+ image = self.center_crop(image, crop_size)
+ if do_rescale:
+ image = self.rescale(image, rescale_factor)
+ if do_normalize:
+ image = self.normalize(image, image_mean, image_std)
+ if do_map_pixels:
+ image = self.map_pixels(image)
+ processed_images.append(image)
+
+ return processed_images
+
+ def _preprocess(
+ self,
+ images: list[np.ndarray],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | None",
+ do_center_crop: bool,
+ crop_size: SizeDict,
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ # Mask related params
+ return_image_mask: bool | None,
+ input_size_patches: int | None,
+ total_mask_patches: int | None,
+ mask_group_min_patches: int | None,
+ mask_group_max_patches: int | None,
+ mask_group_min_aspect_ratio: float | None,
+ mask_group_max_aspect_ratio: float | None,
+ # Codebook related params
+ return_codebook_pixels: bool | None,
+ codebook_do_resize: bool | None,
+ codebook_size: SizeDict | None,
+ codebook_resample: "PILImageResampling | None",
+ codebook_do_center_crop: bool | None,
+ codebook_crop_size: SizeDict | None,
+ codebook_do_rescale: bool | None,
+ codebook_rescale_factor: float | None,
+ codebook_do_map_pixels: bool | None,
+ codebook_do_normalize: bool | None,
+ codebook_image_mean: float | list[float] | None,
+ codebook_image_std: float | list[float] | None,
+ return_tensors: str | TensorType | None,
+ **kwargs,
+ ) -> BatchFeature:
+ processed_images = self._preprocess_image(
+ images=images,
+ do_resize=do_resize,
+ size=size,
+ resample=resample,
+ do_center_crop=do_center_crop,
+ crop_size=crop_size,
+ do_rescale=do_rescale,
+ rescale_factor=rescale_factor,
+ do_normalize=do_normalize,
+ do_map_pixels=False,
+ image_mean=image_mean,
+ image_std=image_std,
+ return_tensors=return_tensors,
+ )
+ data = {
+ "pixel_values": processed_images,
+ }
+
+ if return_codebook_pixels:
+ codebook_processed_images = self._preprocess_image(
+ images=images,
+ do_resize=codebook_do_resize,
+ size=codebook_size,
+ resample=codebook_resample,
+ do_center_crop=codebook_do_center_crop,
+ crop_size=codebook_crop_size,
+ do_rescale=codebook_do_rescale,
+ rescale_factor=codebook_rescale_factor,
+ do_normalize=codebook_do_normalize,
+ do_map_pixels=codebook_do_map_pixels,
+ image_mean=codebook_image_mean,
+ image_std=codebook_image_std,
+ return_tensors=return_tensors,
+ )
+ data["codebook_pixel_values"] = codebook_processed_images
+
+ if return_image_mask:
+ mask_generator = self.masking_generator(
+ input_size_patches=input_size_patches,
+ total_mask_patches=total_mask_patches,
+ mask_group_min_patches=mask_group_min_patches,
+ mask_group_max_patches=mask_group_max_patches,
+ mask_group_min_aspect_ratio=mask_group_min_aspect_ratio,
+ mask_group_max_aspect_ratio=mask_group_max_aspect_ratio,
+ )
+ masks = [mask_generator() for _ in range(len(images))]
+ data["bool_masked_pos"] = masks
+
+ return BatchFeature(data=data, tensor_type=return_tensors)
+
+
+__all__ = ["FlavaImageProcessorPil"]
diff --git a/third_party/transformers/src/transformers/models/flava/modeling_flava.py b/third_party/transformers/src/transformers/models/flava/modeling_flava.py
new file mode 100644
index 0000000000000000000000000000000000000000..49e3b0f3f73caf3a32c5263c2f81c5d7b5eae03a
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/flava/modeling_flava.py
@@ -0,0 +1,1950 @@
+# Copyright 2022 Meta Platforms authors and The HuggingFace Team. All rights reserved.
+#
+# 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 FLAVA model."""
+
+import collections
+import math
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int
+from .configuration_flava import (
+ FlavaConfig,
+ FlavaImageCodebookConfig,
+ FlavaImageConfig,
+ FlavaMultimodalConfig,
+ FlavaTextConfig,
+)
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_CODEBOOK_DOC = "facebook/flava-image-codebook"
+
+LOGIT_SCALE_CLAMP_MIN = 0
+LOGIT_SCALE_CLAMP_MAX = 4.6052
+
+FlavaPossibleConfigs = FlavaTextConfig | FlavaImageConfig | FlavaMultimodalConfig
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Output from FlavaModel containing embeddings and outputs from individual encoders.
+
+ Note that `image_embeddings` and `text_embeddigns` returned are similar to pooled output returned from a
+ transformer. If you want embeddings for contrastive loss or retrieval use a FLAVA model's `image_projection` and
+ `text_projection` layers on `image_embeddings` and `text_embeddings` respectively.
+ """
+)
+class FlavaModelOutput(ModelOutput):
+ r"""
+ image_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `pixel_values` are present):
+ The image embeddings which are basically the pooled output of [`FlavaImageModel`].
+ image_output (`BaseModelOutputWithPooling`, *optional*, returned when `pixel_values` are present):
+ The output of the [`FlavaImageModel`].
+ text_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` are present):
+ The text embeddings which are basically the pooled output of [`FlavaTextModel`].
+ text_output (`BaseModelOutputWithPooling`, *optional*, returned when `input_ids` are present):
+ The output of the [`FlavaTextModel`].
+ multimodal_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` and `pixel_values` are present and `skip_multimodal_encoder` is `None` or `False`):
+ The multimodal embeddings which are basically the pooled output of [`FlavaTextModel`].
+ multimodal_output (`BaseModelOutputWithPooling`, returned when `input_ids` and `pixel_values` are present and `skip_multimodal_encoder` is `None` or `False`):
+ The output of the [`FlavaMultimodalModel`].
+ """
+
+ image_embeddings: torch.FloatTensor | None = None
+ image_output: BaseModelOutputWithPooling | None = None
+ text_embeddings: torch.FloatTensor | None = None
+ text_output: BaseModelOutputWithPooling | None = None
+ multimodal_embeddings: torch.FloatTensor | None = None
+ multimodal_output: BaseModelOutputWithPooling | None = None
+
+ def to_tuple(self) -> tuple[Any]:
+ return tuple(
+ self[k] if k not in ["text_output", "image_output", "multimodal_output"] else getattr(self, k).to_tuple()
+ for k in self.keys()
+ )
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Class representing pretraining losses from FLAVA model
+ """
+)
+class FlavaLosses(ModelOutput):
+ r"""
+ mim (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mim_labels` and `pixel_values` are present, `input_ids_masked` is absent and `mim_weight` > 0.):
+ Masked Image Modeling loss as used in BeIT calculated only for unimodal image data.
+ mlm (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mlm_labels` and `input_ids_masked` are present, `pixel_values` is absent and `mlm_weight` > 0.):
+ Masked Language Modeling loss as used in BERT calculated only for unimodal text data.
+ itm (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `itm_labels`, `input_ids_masked`, `pixel_values` are present and `itm_weight` > 0.):
+ Image Text Matching (ITM) loss calculated for paired image-text data. Note that ITM loss is calculated on
+ masked pairs in FLAVA.
+ global_contrastive (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `input_ids` and `pixel_values` are present and `global_contrastive_weight` > 0.):
+ Contrastive loss for image-text similarity similar to CLIP but calculated globally for paired image-text
+ data. This is calculated on unmasked images and texts.
+ mmm_image (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mim_labels`, `pixel_values` and `input_ids_masked` are present and `mmm_image_weight` > 0.):
+ Masked Multimodal Modeling loss's image component calculated on paired image-text data.
+ mmm_text (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `mlm_labels`, `pixel_values` and `input_ids_masked` are present and `mmm_text_weight` > 0.):
+ Masked Multimodal Modeling loss's text component calculated on paired image-text data.
+ """
+
+ mim: torch.FloatTensor | None = None
+ mlm: torch.FloatTensor | None = None
+ itm: torch.FloatTensor | None = None
+ global_contrastive: torch.FloatTensor | None = None
+ mmm_image: torch.FloatTensor | None = None
+ mmm_text: torch.FloatTensor | None = None
+
+ def all_none(self) -> bool:
+ all_none = True
+ for v in self.values():
+ if v is not None:
+ all_none = False
+ break
+ return all_none
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Output from FlavaForPreTraining containing embeddings, and outputs from individual encoders.
+
+ Note that `image_embeddings` and `text_embeddings` returned are similar to pooled output returned from a
+ transformer. If you want embeddings for contrastive loss or retrieval use a FLAVA model's `image_projection` and
+ `text_projection` layers on `image_embeddings` and `text_embeddings` respectively.
+ """
+)
+class FlavaForPreTrainingOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor`, *optional*, returned when `return_loss` is True):
+ Total loss calculated for this model.
+ loss_info (`FlavaLosses`):
+ Detailed info for FLAVA Pretraining losses. Check `FlavaLosses` class description for the information on
+ the keys.
+ image_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `pixel_values` are present):
+ The image embeddings which are basically the pooled output of [`FlavaImageModel`].
+ image_output (`BaseModelOutputWithPooling`, *optional*, returned when `pixel_values` are present):
+ The output of the [`FlavaImageModel`].
+ text_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` are present):
+ The text embeddings which are basically the pooled output of [`FlavaTextModel`].
+ text_output (`BaseModelOutputWithPooling`, *optional*, returned when `input_ids` are present):
+ The output of the [`FlavaTextModel`].
+ multimodal_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` and `pixel_values` are present and `skip_unmasked_multimodal_encoder` is `None` or `False`):
+ The multimodal embeddings which are basically the pooled output of [`FlavaTextModel`].
+ multimodal_output (`BaseModelOutputWithPooling`, returned when `input_ids` and `pixel_values` are present and `skip_unmasked_multimodal_encoder` is `None` or `False`):
+ The output of the [`FlavaMultimodalModel`].
+ image_masked_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `pixel_values` are present):
+ The image embeddings which are basically the pooled output of [`FlavaImageModel`]. Uses `bool_masked_pos`
+ to create masked images.
+ image_masked_output (`BaseModelOutputWithPooling`, *optional*, returned when `pixel_values` are present):
+ The output of the [`FlavaImageModel`]. Uses `bool_masked_pos` to create masked images.
+ text_masked_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids_masked` are present):
+ The text embeddings which are basically the pooled output of [`FlavaTextModel`].
+ text_masked_output (`BaseModelOutputWithPooling`, *optional*, returned when `input_ids_masked` are present):
+ The output of the [`FlavaTextModel`].
+ multimodal_masked_embeddings (`torch.FloatTensor` of shape `(batch_size, output_dim)`, *optional*, returned when `input_ids` and `pixel_values` are present):
+ The multimodal embeddings which are basically the pooled output of [`FlavaTextModel`].
+ multimodal_masked_output (`BaseModelOutputWithPooling`, *optional*, returned when `input_ids_masked` and `pixel_values` are present):
+ The output of the [`FlavaMultimodalModel`].
+ mim_logits (`torch.FloatTensor` of shape `(batch_size, num_image_patches, image_vocab_size)` or of shape `(total_masked_patches, image_vocab_size)` , *optional*, returned when `pixel_values` are present and `input_ids_masked` are not):
+ The logits for MIM unimodal loss. Uses `book_masked_pos` to get masked patches. The flattened output is
+ returned when `bool_masked_pos` has some of the patches masked.
+ mlm_logits (`torch.FloatTensor` of shape `(batch_size, text_seq_length, text_vocab_size)` or of shape `(total_masked_seq_length, text_vocab_size)`, *optional*, returned when `input_ids_masked` are present and `pixel_values` are not):
+ The logits for MLM unimodal loss. The flattened output is returned when `input_ids_masked` has some of
+ the tokens masked.
+ itm_logits (`torch.FloatTensor` of shape `(batch_size, 2)`, *optional*, returned when `input_ids_masked` and `pixel_values` are present):
+ The logits for ITM loss. Note that ITM loss is calculated on masked pairs in FLAVA.
+ contrastive_logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`):
+ The scaled dot product scores between `image_embeddings` and `text_embeddings` but passed through FLAVA's
+ `image_projection` and `text_projection` layers respectively. This represents the image-text similarity
+ scores. This is calculated on unmasked images and texts.
+ contrastive_logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`):
+ The scaled dot product scores between `text_embeddings` and `image_embeddings` but passed through FLAVA's
+ `text_projection` and `image_projection` layers respectively. This is calculated on unmasked images and
+ texts.
+ mmm_image_logits (`torch.FloatTensor` of shape `(batch_size, num_image_patches, image_vocab_size)` or of shape`(total_masked_patches, image_vocab_size)`, *optional*, returned when `pixel_values` and `input_ids_masked` are present):
+ The logits for MMM image multimodal loss. Uses `book_masked_pos` to get masked patches. The flattened
+ output is returned when `bool_masked_pos` has some of the patches masked.
+ mmm_text_logits (`torch.FloatTensor` of shape `(batch_size, text_seq_length, text_vocab_size)` or of shape `(`(total_masked_seq_length, text_vocab_size)`), *optional*, returned when `pixel_values` and `input_ids_masked` are present):
+ The logits for MMM text multimodal loss. The flattened output is returned when `input_ids_masked` has
+ some of the tokens masked.
+ """
+
+ loss: torch.FloatTensor | None = None
+ loss_info: FlavaLosses = None
+ image_embeddings: torch.FloatTensor | None = None
+ image_output: BaseModelOutputWithPooling | None = None
+ text_embeddings: torch.FloatTensor | None = None
+ text_output: BaseModelOutputWithPooling | None = None
+ multimodal_embeddings: torch.FloatTensor | None = None
+ multimodal_output: BaseModelOutputWithPooling | None = None
+ image_masked_embeddings: torch.FloatTensor | None = None
+ image_masked_output: BaseModelOutputWithPooling | None = None
+ text_masked_embeddings: torch.FloatTensor | None = None
+ text_masked_output: BaseModelOutputWithPooling | None = None
+ multimodal_masked_embeddings: torch.FloatTensor | None = None
+ multimodal_masked_output: BaseModelOutputWithPooling | None = None
+ mim_logits: torch.FloatTensor | None = None
+ mlm_logits: torch.FloatTensor | None = None
+ itm_logits: torch.FloatTensor | None = None
+ contrastive_logits_per_image: torch.FloatTensor | None = None
+ contrastive_logits_per_text: torch.FloatTensor | None = None
+ mmm_image_logits: torch.FloatTensor | None = None
+ mmm_text_logits: torch.FloatTensor | None = None
+
+ def to_tuple(self) -> tuple[Any]:
+ transformer_outputs = [
+ "text_output",
+ "image_output",
+ "multimodal_output",
+ "text_masked_output",
+ "image_masked_output",
+ "multimodal_masked_output",
+ ]
+ return tuple(self[k] if k not in transformer_outputs else getattr(self, k).to_tuple() for k in self.keys())
+
+
+# Based on timm implementation, which can be found here:
+# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/image_transformer.py
+class FlavaImageEmbeddings(nn.Module):
+ """
+ Construct the CLS token, position and patch embeddings. Optionally, also the mask token.
+ """
+
+ def __init__(self, config: FlavaImageConfig, use_mask_token: bool = False) -> None:
+ super().__init__()
+
+ use_mask_token = use_mask_token or config.mask_token
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
+ self.mask_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) if use_mask_token else None
+ self.patch_embeddings = PatchEmbeddings(
+ image_size=config.image_size,
+ patch_size=config.patch_size,
+ num_channels=config.num_channels,
+ embed_dim=config.hidden_size,
+ )
+ num_patches = self.patch_embeddings.num_patches
+ self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.hidden_size))
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ self.patch_size = config.patch_size
+ self.config = config
+
+ # Copied from transformers.models.vit.modeling_vit.ViTEmbeddings.interpolate_pos_encoding
+ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+ """
+ This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution
+ images. This method is also adapted to support torch.jit tracing.
+
+ Adapted from:
+ - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
+ - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
+ """
+
+ num_patches = embeddings.shape[1] - 1
+ num_positions = self.position_embeddings.shape[1] - 1
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embeddings
+
+ class_pos_embed = self.position_embeddings[:, :1]
+ patch_pos_embed = self.position_embeddings[:, 1:]
+
+ dim = embeddings.shape[-1]
+
+ new_height = height // self.patch_size
+ new_width = width // self.patch_size
+
+ sqrt_num_positions = torch_int(num_positions**0.5)
+ patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
+ patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed,
+ size=(new_height, new_width),
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+ return torch.cat((class_pos_embed, patch_pos_embed), dim=1)
+
+ def forward(
+ self,
+ pixel_values: torch.Tensor,
+ bool_masked_pos: torch.BoolTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ ) -> torch.Tensor:
+ batch_size, num_channels, height, width = pixel_values.shape
+ embeddings = self.patch_embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
+
+ batch_size, seq_len, _ = embeddings.size()
+ if bool_masked_pos is not None:
+ mask_tokens = self.mask_token.expand(batch_size, seq_len, -1)
+ # B X H X W = B X HW
+ if bool_masked_pos.dim() == 3:
+ bool_masked_pos = bool_masked_pos.view(bool_masked_pos.size(0), -1)
+ # replace the masked visual tokens by mask_tokens
+ mask = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens)
+ embeddings = embeddings * (1.0 - mask) + mask_tokens * mask
+
+ # add the [CLS] token to the embedded patch tokens
+ cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+ embeddings = torch.cat((cls_tokens, embeddings), dim=1)
+
+ # add positional encoding to each token
+ if interpolate_pos_encoding:
+ embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+ else:
+ embeddings = embeddings + self.position_embeddings
+
+ embeddings = self.dropout(embeddings)
+
+ return embeddings
+
+
+# Based on timm implementation, which can be found here:
+# https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/image_transformer.py
+class PatchEmbeddings(nn.Module):
+ """
+ Image to Patch Embedding.
+ """
+
+ def __init__(
+ self,
+ image_size: int | list[int] | tuple[int, int] = 224,
+ patch_size: int | tuple[int, int] = 16,
+ num_channels: int = 3,
+ embed_dim: int = 768,
+ ):
+ super().__init__()
+ if not isinstance(image_size, collections.abc.Iterable):
+ image_size = (image_size, image_size)
+ if not isinstance(patch_size, collections.abc.Iterable):
+ patch_size = (patch_size, patch_size)
+ num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
+ self.image_size = image_size
+ self.patch_size = patch_size
+ self.num_patches = num_patches
+
+ self.projection = nn.Conv2d(num_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+ def forward(self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False) -> torch.Tensor:
+ batch_size, num_channels, height, width = pixel_values.shape
+ if not interpolate_pos_encoding:
+ if height != self.image_size[0] or width != self.image_size[1]:
+ raise ValueError(
+ f"Input image size ({height}*{width}) doesn't match model"
+ f" ({self.image_size[0]}*{self.image_size[1]})."
+ )
+ x = self.projection(pixel_values).flatten(2).transpose(1, 2)
+ return x
+
+
+class FlavaTextEmbeddings(nn.Module):
+ """Construct the embeddings from word, position and token_type embeddings."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+ self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+ self.register_buffer(
+ "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+ )
+ self.register_buffer(
+ "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
+ )
+
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ ):
+ input_shape = input_ids.size()
+ seq_length = input_shape[1]
+
+ if position_ids is None:
+ position_ids = self.position_ids[:, :seq_length]
+
+ # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
+ # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
+ # issue #5664
+ if token_type_ids is None:
+ if hasattr(self, "token_type_ids"):
+ buffered_token_type_ids = self.token_type_ids[:, :seq_length]
+ buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
+ token_type_ids = buffered_token_type_ids_expanded
+ else:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+
+ inputs_embeds = self.word_embeddings(input_ids)
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
+ embeddings = inputs_embeds + token_type_embeddings
+
+ position_embeddings = self.position_embeddings(position_ids)
+ embeddings += position_embeddings
+
+ embeddings = self.LayerNorm(embeddings)
+ embeddings = self.dropout(embeddings)
+ return embeddings
+
+
+class FlavaSelfAttention(nn.Module):
+ def __init__(self, config: FlavaPossibleConfigs) -> None:
+ super().__init__()
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+ raise ValueError(
+ f"The hidden size {config.hidden_size} is not a multiple of the number of attention "
+ f"heads {config.num_attention_heads}."
+ )
+
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
+ self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ output_attentions: bool = False,
+ ) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor]:
+ batch_size, seq_length, _ = hidden_states.shape
+ query_layer = (
+ self.query(hidden_states)
+ .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)
+ .transpose(1, 2)
+ )
+ key_layer = (
+ self.key(hidden_states)
+ .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)
+ .transpose(1, 2)
+ )
+ value_layer = (
+ self.value(hidden_states)
+ .view(batch_size, -1, self.num_attention_heads, self.attention_head_size)
+ .transpose(1, 2)
+ )
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+ attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+ if attention_mask is not None:
+ # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
+ attention_scores = attention_scores + attention_mask
+
+ # Normalize the attention scores to probabilities.
+ attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+
+ # This is actually dropping out entire tokens to attend to, which might
+ # seem a bit unusual, but is taken from the original Transformer paper.
+ attention_probs = self.dropout(attention_probs)
+
+ context_layer = torch.matmul(attention_probs, value_layer)
+
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+ context_layer = context_layer.view(*new_context_layer_shape)
+
+ outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+
+ return outputs
+
+
+class FlavaSelfOutput(nn.Module):
+ """
+ The residual connection is defined in FlavaLayer (same as ViTLayer) instead of here (as is the case with other
+ models), due to the layernorm applied before each block.
+ """
+
+ def __init__(self, config: FlavaPossibleConfigs) -> None:
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+
+ return hidden_states
+
+
+class FlavaAttention(nn.Module):
+ def __init__(self, config: FlavaPossibleConfigs) -> None:
+ super().__init__()
+ self.attention = FlavaSelfAttention(config)
+ self.output = FlavaSelfOutput(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ output_attentions: bool = False,
+ ) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor]:
+ self_outputs = self.attention(
+ hidden_states, attention_mask=attention_mask, output_attentions=output_attentions
+ )
+
+ attention_output = self.output(self_outputs[0], hidden_states)
+
+ outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
+ return outputs
+
+
+class FlavaIntermediate(nn.Module):
+ def __init__(self, config: FlavaPossibleConfigs) -> None:
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ # Copied from transformers.models.vit.modeling_vit.ViTIntermediate.forward
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+
+ return hidden_states
+
+
+class FlavaOutput(nn.Module):
+ def __init__(self, config: FlavaPossibleConfigs) -> None:
+ super().__init__()
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ # Copied from transformers.models.vit.modeling_vit.ViTOutput.forward
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+
+ hidden_states = hidden_states + input_tensor
+
+ return hidden_states
+
+
+class FlavaLayer(GradientCheckpointingLayer):
+ """This corresponds to the Block class in the timm implementation."""
+
+ def __init__(self, config: FlavaPossibleConfigs) -> None:
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+ self.attention = FlavaAttention(config)
+ self.intermediate = FlavaIntermediate(config)
+ self.output = FlavaOutput(config)
+
+ # TODO: Check fp32 layer norm possibility
+ self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ output_attentions: bool = False,
+ ) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor]:
+ self_attention_outputs = self.attention(
+ self.layernorm_before(hidden_states), # in ViT, layernorm is applied before self-attention
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ )
+ attention_output = self_attention_outputs[0]
+ outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
+
+ # first residual connection
+ hidden_states = attention_output + hidden_states
+
+ # in ViT, layernorm is also applied after self-attention
+ layer_output = self.layernorm_after(hidden_states)
+ layer_output = self.intermediate(layer_output)
+
+ # second residual connection is done here
+ layer_output = self.output(layer_output, hidden_states)
+
+ outputs = (layer_output,) + outputs
+
+ return outputs
+
+
+class FlavaEncoder(nn.Module):
+ def __init__(self, config: FlavaConfig) -> None:
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList([FlavaLayer(config) for _ in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ output_attentions: bool = False,
+ output_hidden_states: bool = False,
+ return_dict: bool = True,
+ ) -> tuple | BaseModelOutput:
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attentions = () if output_attentions else None
+
+ for i, layer_module in enumerate(self.layer):
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ layer_outputs = layer_module(hidden_states, attention_mask, output_attentions)
+
+ hidden_states = layer_outputs[0]
+
+ if output_attentions:
+ all_self_attentions = all_self_attentions + (layer_outputs[1],)
+
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ if not return_dict:
+ return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
+ return BaseModelOutput(
+ last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_self_attentions
+ )
+
+
+class FlavaPooler(nn.Module):
+ def __init__(self, config: FlavaPossibleConfigs):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.activation = nn.Tanh()
+
+ def forward(self, hidden_states: torch.Tensor):
+ # We "pool" the model by simply taking the hidden state corresponding
+ # to the first token.
+ first_token_tensor = hidden_states[:, 0]
+ pooled_output = self.dense(first_token_tensor)
+ pooled_output = self.activation(pooled_output)
+ return pooled_output
+
+
+@auto_docstring
+class FlavaPreTrainedModel(PreTrainedModel):
+ config: FlavaConfig
+ base_model_prefix = "flava"
+ input_modalities = ("image", "text")
+ supports_gradient_checkpointing = True
+
+ @torch.no_grad()
+ def _init_weights(self, module: nn.Linear | nn.Conv2d | nn.LayerNorm) -> None:
+ """Initialize the weights"""
+ super()._init_weights(module)
+ if isinstance(module, FlavaMaskedPredictionHead):
+ init.zeros_(module.bias)
+ elif isinstance(module, FlavaImageEmbeddings):
+ init.zeros_(module.cls_token)
+ init.zeros_(module.position_embeddings)
+ if module.mask_token is not None:
+ init.zeros_(module.mask_token)
+ elif isinstance(module, FlavaTextEmbeddings):
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+ init.zeros_(module.token_type_ids)
+ elif isinstance(module, FlavaMultimodalModel):
+ if module.use_cls_token:
+ init.zeros_(module.cls_token)
+ elif isinstance(module, FlavaModel):
+ init.constant_(module.logit_scale, self.config.logit_scale_init_value)
+
+
+@auto_docstring
+class FlavaImageModel(FlavaPreTrainedModel):
+ config: FlavaImageConfig
+ # This override allows us to load FlavaImageModel from FlavaModel/FlavaForPreTraining checkpoints.
+ base_model_prefix = "flava.image_model"
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+
+ def __init__(self, config: FlavaImageConfig, add_pooling_layer: bool = True):
+ r"""
+ add_pooling_layer (bool, *optional*, defaults to `True`):
+ Whether to add a pooling layer
+ """
+ super().__init__(config)
+
+ self.config = config
+
+ self.embeddings = FlavaImageEmbeddings(config)
+ self.encoder = FlavaEncoder(config)
+
+ self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.pooler = FlavaPooler(config) if add_pooling_layer else None
+
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.embeddings.patch_embeddings
+
+ def set_input_embeddings(self, value: nn.Module):
+ self.embeddings.patch_embeddings = value
+
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.Tensor | None = None,
+ bool_masked_pos: torch.BoolTensor | None = None,
+ interpolate_pos_encoding: bool | None = None,
+ attention_mask: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, image_num_patches)`):
+ Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
+ """
+ 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.return_dict
+
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ embedding_output = self.embeddings(
+ pixel_values, bool_masked_pos=bool_masked_pos, interpolate_pos_encoding=interpolate_pos_encoding
+ )
+
+ encoder_outputs = self.encoder(
+ embedding_output,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ sequence_output = encoder_outputs[0]
+ sequence_output = self.layernorm(sequence_output)
+ pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+ if not return_dict:
+ return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ )
+
+
+@auto_docstring
+class FlavaTextModel(FlavaPreTrainedModel):
+ config: FlavaTextConfig
+ # This override allows us to load FlavaTextModel from FlavaModel/FlavaForPreTraining checkpoints.
+ base_model_prefix = "flava.text_model"
+ input_modalities = ("text",)
+
+ def __init__(self, config: FlavaTextConfig, add_pooling_layer: bool = True):
+ r"""
+ add_pooling_layer (bool, *optional*, defaults to `True`):
+ Whether to add a pooling layer
+ """
+ super().__init__(config)
+ self.config = config
+
+ self.embeddings = FlavaTextEmbeddings(config)
+ self.encoder = FlavaEncoder(config)
+
+ self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.pooler = FlavaPooler(config) if add_pooling_layer else None
+
+ self.post_init()
+
+ def get_input_embeddings(self) -> PatchEmbeddings:
+ return self.embeddings.word_embeddings
+
+ def set_input_embeddings(self, value: nn.Module):
+ self.embeddings.word_embeddings = value
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, text_seq_length)`):
+ Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See
+ [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input
+ IDs?](../glossary#input-ids)
+ token_type_ids (`torch.LongTensor` of shape `(batch_size, text_seq_length)`, *optional*):
+ Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
+ 1]`:
+ - 0 corresponds to a *sentence A* token,
+ - 1 corresponds to a *sentence B* token.
+ [What are token type IDs?](../glossary#token-type-ids)
+ """
+ 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.return_dict
+
+ if input_ids is None:
+ raise ValueError("You have to specify input_ids")
+
+ input_shape = input_ids.size()
+
+ if attention_mask is None:
+ attention_mask = torch.ones(input_shape, device=input_ids.device)
+
+ extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(
+ attention_mask,
+ input_shape,
+ )
+
+ embedding_output = self.embeddings(
+ input_ids=input_ids,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ )
+
+ encoder_outputs = self.encoder(
+ embedding_output,
+ attention_mask=extended_attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ sequence_output = encoder_outputs[0]
+ sequence_output = self.layernorm(sequence_output)
+ pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+ if not return_dict:
+ return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ )
+
+
+@auto_docstring
+class FlavaMultimodalModel(FlavaPreTrainedModel):
+ config: FlavaMultimodalConfig
+ # This override allows us to load FlavaMultimodalModel from FlavaModel/FlavaForPreTraining checkpoints.
+ base_model_prefix = "flava.multimodal_model"
+ main_input_name = "hidden_states"
+
+ def __init__(self, config: FlavaMultimodalConfig, add_pooling_layer=True):
+ r"""
+ add_pooling_layer (bool, *optional*, defaults to `True`):
+ Whether to add a pooling layer
+ """
+ super().__init__(config)
+ self.config = config
+ self.use_cls_token = self.config.use_cls_token
+ if self.use_cls_token:
+ self.cls_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size))
+
+ self.encoder = FlavaEncoder(config)
+
+ self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.pooler = FlavaPooler(config) if add_pooling_layer else None
+
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ hidden_states (`torch.FloatTensor` of shape `(batch_size, image_num_patches + text_seq_len, hidden_size)`):
+ The concatenated hidden states of unimodal encoders.
+ """
+ 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.return_dict
+
+ batch_size, seq_length, _ = hidden_states.size()
+
+ if self.use_cls_token:
+ cls_tokens = self.cls_token.expand(batch_size, -1, -1)
+ hidden_states = torch.cat((cls_tokens, hidden_states), dim=1)
+ seq_length += 1
+
+ if attention_mask is None:
+ attention_mask = torch.ones((batch_size, seq_length), device=hidden_states.device)
+
+ extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(
+ attention_mask,
+ (batch_size, seq_length),
+ )
+
+ encoder_outputs = self.encoder(
+ hidden_states,
+ attention_mask=extended_attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ sequence_output = encoder_outputs[0]
+ sequence_output = self.layernorm(sequence_output)
+ pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+ if not return_dict:
+ return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ )
+
+
+@auto_docstring
+class FlavaModel(FlavaPreTrainedModel):
+ config: FlavaConfig
+
+ def __init__(self, config: FlavaConfig):
+ super().__init__(config)
+
+ if not isinstance(config.text_config, FlavaTextConfig):
+ raise TypeError(
+ "config.text_config is expected to be of type FlavaTextConfig but is of type"
+ f" {type(config.text_config)}."
+ )
+
+ if not isinstance(config.image_config, FlavaImageConfig):
+ raise TypeError(
+ "config.image_config is expected to be of type FlavaImageConfig but is of type"
+ f" {type(config.image_config)}."
+ )
+
+ if not isinstance(config.multimodal_config, FlavaMultimodalConfig):
+ raise TypeError(
+ "config.multimodal_config is expected to be of type FlavaMultimodalConfig but "
+ + f"is of type {type(config.multimodal_config)}."
+ )
+
+ text_config = config.text_config
+ image_config = config.image_config
+ multimodal_config = config.multimodal_config
+
+ self.projection_dim = config.projection_dim
+ self.text_hidden_size = text_config.hidden_size
+ self.image_hidden_size = image_config.hidden_size
+ self.mm_hidden_size = multimodal_config.hidden_size
+
+ self.text_model = FlavaTextModel(text_config)
+ self.image_model = FlavaImageModel(image_config)
+ self.multimodal_model = FlavaMultimodalModel(multimodal_config)
+
+ self.image_projection = nn.Linear(self.image_hidden_size, self.projection_dim)
+ self.text_projection = nn.Linear(self.text_hidden_size, self.projection_dim)
+ self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))
+
+ self.image_to_mm_projection = nn.Linear(self.image_hidden_size, self.mm_hidden_size)
+ self.text_to_mm_projection = nn.Linear(self.text_hidden_size, self.mm_hidden_size)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def get_text_features(
+ self,
+ input_ids: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, text_seq_length)`):
+ Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See
+ [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input
+ IDs?](../glossary#input-ids)
+ token_type_ids (`torch.LongTensor` of shape `(batch_size, text_seq_length)`, *optional*):
+ Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
+ 1]`:
+ - 0 corresponds to a *sentence A* token,
+ - 1 corresponds to a *sentence B* token.
+ [What are token type IDs?](../glossary#token-type-ids)
+
+ Examples:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoProcessor, FlavaModel
+
+ >>> model = FlavaModel.from_pretrained("{0}")
+ >>> processor = AutoProcessor.from_pretrained("{0}")
+
+ >>> inputs = processor(
+ ... text=["a photo of a cat", "a photo of a dog"], max_length=77, padding="max_length", return_tensors="pt"
+ ... )
+ >>> with torch.inference_mode():
+ ... text_features = model.get_text_features(**inputs)
+ ```
+ """
+ text_outputs: BaseModelOutputWithPooling = self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ return_dict=True,
+ **kwargs,
+ )
+ last_hidden_state = text_outputs.last_hidden_state
+ text_outputs.pooler_output = self.text_projection(last_hidden_state)
+
+ return text_outputs
+
+ @can_return_tuple
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.Tensor,
+ bool_masked_pos: torch.BoolTensor | None = None,
+ interpolate_pos_encoding: bool | None = None,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, image_num_patches)`):
+ Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
+
+ Examples:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoProcessor, FlavaModel
+ >>> from transformers.image_utils import load_image
+
+ >>> model = FlavaModel.from_pretrained("{0}")
+ >>> processor = AutoProcessor.from_pretrained("{0}")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = load_image(url)
+
+ >>> inputs = processor(images=image, return_tensors="pt")
+
+ >>> with torch.inference_mode():
+ ... image_features = model.get_image_features(**inputs)
+ ```
+ """
+ image_outputs: BaseModelOutputWithPooling = self.image_model(
+ pixel_values=pixel_values,
+ bool_masked_pos=bool_masked_pos,
+ attention_mask=attention_mask,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ return_dict=True,
+ **kwargs,
+ )
+ last_hidden_state = image_outputs.last_hidden_state
+ image_outputs.pooler_output = self.image_projection(last_hidden_state)
+
+ return image_outputs
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ bool_masked_pos: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ image_attention_mask: torch.Tensor | None = None,
+ skip_multimodal_encoder: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool = True,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | FlavaModelOutput:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, image_num_patches + text_seq_len)`):
+ Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See
+ [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input
+ IDs?](../glossary#input-ids)
+ token_type_ids (`torch.LongTensor` of shape `(batch_size, image_num_patches + text_seq_len)`, *optional*):
+ Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
+ 1]`:
+ - 0 corresponds to a *sentence A* token,
+ - 1 corresponds to a *sentence B* token.
+ [What are token type IDs?](../glossary#token-type-ids)
+ bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, image_num_patches)`):
+ Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
+ image_attention_mask (`torch.Tensor` of shape `(batch_size, image_num_patches)`, *optional*):
+ Mask to avoid performing attention on padding pixel values for image inputs. Mask values selected in `[0, 1]`:
+ - 1 for pixel values that are real (i.e., **not masked**),
+ - 0 for pixel values that are padding (i.e., **masked**).
+ skip_multimodal_encoder (*bool*, *optional*):
+ Skip any calculations for multimodal encoder. Useful if multimodal encoding is not going to be used.
+
+ Examples:
+
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoProcessor, FlavaModel
+
+ >>> model = FlavaModel.from_pretrained("facebook/flava-full")
+ >>> processor = AutoProcessor.from_pretrained("facebook/flava-full")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> inputs = processor(text=["a photo of a cat"], images=image, return_tensors="pt", padding=True)
+
+ >>> outputs = model(**inputs)
+
+ >>> image_embeddings = outputs.image_embeddings
+ >>> text_embeddings = outputs.text_embeddings
+ >>> multimodal_embeddings = outputs.multimodal_embeddings
+
+ >>> outputs.image_embeddings.shape
+ torch.Size([1, 197, 768])
+
+ >>> text_embeddings.shape
+ torch.Size([1, 7, 768])
+
+ >>> multimodal_embeddings.shape
+ torch.Size([1, 205, 768])
+ ```
+ """
+
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+ if not output_hidden_states:
+ raise ValueError("FLAVA model requires hidden states to work. Please set `output_hidden_states=True`")
+ image_embeddings = None
+ image_states = None
+ image_mm_projection = None
+ image_output = None
+ if pixel_values is not None:
+ image_output = self.image_model(
+ pixel_values=pixel_values,
+ bool_masked_pos=bool_masked_pos,
+ attention_mask=image_attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ image_embeddings, image_states = image_output[0], image_output[2]
+ # Note that these states don't use final layernorm in the transformer model
+ image_mm_projection = self.image_to_mm_projection(image_states[-1])
+
+ text_embeddings = None
+ text_states = None
+ text_mm_projection = None
+ text_output = None
+ if input_ids is not None:
+ text_output = self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ token_type_ids=token_type_ids,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ text_embeddings, text_states = text_output[0], text_output[2]
+ # Note that these states don't use final layernorm in the transformer model
+ text_mm_projection = self.text_to_mm_projection(text_states[-1])
+
+ multimodal_embeddings = None
+ multimodal_output = None
+ if image_mm_projection is not None and text_mm_projection is not None and not skip_multimodal_encoder:
+ if attention_mask is not None:
+ batch_size, seq_len, _ = image_mm_projection.shape
+ if self.multimodal_model.use_cls_token:
+ seq_len += 1
+ attention_mask_image = torch.ones(batch_size, seq_len, device=image_mm_projection.device)
+ attention_multimodal = torch.cat([attention_mask_image, attention_mask], dim=1)
+ else:
+ attention_multimodal = None
+ multimodal_input = torch.cat([image_mm_projection, text_mm_projection], dim=1)
+ multimodal_output = self.multimodal_model(
+ multimodal_input, attention_mask=attention_multimodal, return_dict=return_dict
+ )
+ multimodal_embeddings = multimodal_output[0]
+
+ if not return_dict:
+ return (
+ image_embeddings,
+ image_output,
+ text_embeddings,
+ text_output,
+ multimodal_embeddings,
+ multimodal_output,
+ )
+
+ return FlavaModelOutput(
+ image_embeddings=image_embeddings,
+ image_output=image_output,
+ text_embeddings=text_embeddings,
+ text_output=text_output,
+ multimodal_embeddings=multimodal_embeddings,
+ multimodal_output=multimodal_output,
+ )
+
+
+class FlavaImageCodebookResPath(nn.Module):
+ def __init__(self, in_size: int, out_size: int, **kwargs):
+ super().__init__()
+ hid_size = out_size // 4
+
+ path = OrderedDict()
+ path["relu_1"] = nn.ReLU()
+ path["conv_1"] = nn.Conv2d(in_size, hid_size, kernel_size=3, padding=1)
+ path["relu_2"] = nn.ReLU()
+ path["conv_2"] = nn.Conv2d(hid_size, hid_size, kernel_size=3, padding=1)
+ path["relu_3"] = nn.ReLU()
+ path["conv_3"] = nn.Conv2d(hid_size, hid_size, kernel_size=3, padding=1)
+ path["relu_4"] = nn.ReLU()
+ path["conv_4"] = nn.Conv2d(hid_size, out_size, kernel_size=1, padding=0)
+
+ self.path = nn.Sequential(path)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.path(x)
+
+
+class FlavaImageCodebookBlock(nn.Module):
+ def __init__(self, in_size: int, out_size: int, num_layers: int, **kwargs):
+ super().__init__()
+
+ self.post_gain = 1 / (num_layers**2)
+
+ if in_size != out_size:
+ self.id_path = nn.Conv2d(in_size, out_size, kernel_size=1, padding=0)
+ else:
+ self.id_path = nn.Identity()
+
+ self.res_path = FlavaImageCodebookResPath(in_size, out_size)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.id_path(x) + self.post_gain * self.res_path(x)
+
+
+class FlavaImageCodebookLayerGroup(nn.Module):
+ def __init__(self, num_blocks: int, num_layers: int, in_size: int, out_size: int, use_pool: bool = True):
+ super().__init__()
+ blocks = OrderedDict()
+ for i in range(num_blocks):
+ if i == 0:
+ blocks[f"block_{i + 1}"] = FlavaImageCodebookBlock(in_size, out_size, num_layers)
+ else:
+ blocks[f"block_{i + 1}"] = FlavaImageCodebookBlock(out_size, out_size, num_layers)
+
+ if use_pool:
+ blocks["pool"] = nn.MaxPool2d(kernel_size=2)
+
+ self.group = nn.Sequential(blocks)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.group(x)
+
+
+# Inspired by DALLE Encoder in https://github.com/openai/DALL-E/blob/5be4b236bc3ade6943662354117a0e83752cc322/dall_e/encoder.py#L42
+@auto_docstring(
+ custom_intro="""
+ The FLAVA's image codebook model inspired from DALL-E's original encoder. Outputs raw hidden states and can be used
+ to generate image tokens for an image based on DALL-E's vocab. Used to generate labels for MIM. Use
+ `get_codebook_indices` to get image tokens for an image.
+ """
+)
+class FlavaImageCodebook(FlavaPreTrainedModel):
+ base_model_prefix = "model"
+ config: FlavaImageCodebookConfig
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+ supports_gradient_checkpointing = False
+
+ def __init__(
+ self,
+ config: FlavaImageCodebookConfig,
+ **kwargs: Any,
+ ):
+ super().__init__(config)
+
+ self.config = config
+ self.num_groups = config.num_groups
+ self.input_channels = config.input_channels
+ self.num_blocks_per_group = config.num_blocks_per_group
+ self.hidden_size = config.hidden_size
+ self.vocab_size = config.vocab_size
+
+ num_layers = self.num_groups * self.num_blocks_per_group
+
+ output_blocks = OrderedDict()
+ output_blocks["relu"] = nn.ReLU()
+ output_blocks["conv"] = nn.Conv2d(8 * self.hidden_size, self.vocab_size, kernel_size=1, padding=0)
+
+ blocks = OrderedDict()
+ blocks["input"] = nn.Conv2d(self.input_channels, 1 * self.hidden_size, kernel_size=7, padding=3)
+ blocks["group_1"] = FlavaImageCodebookLayerGroup(
+ self.num_blocks_per_group, num_layers, 1 * self.hidden_size, 1 * self.hidden_size
+ )
+ blocks["group_2"] = FlavaImageCodebookLayerGroup(
+ self.num_blocks_per_group, num_layers, 1 * self.hidden_size, 2 * self.hidden_size
+ )
+ blocks["group_3"] = FlavaImageCodebookLayerGroup(
+ self.num_blocks_per_group, num_layers, 2 * self.hidden_size, 4 * self.hidden_size
+ )
+ blocks["group_4"] = FlavaImageCodebookLayerGroup(
+ self.num_blocks_per_group, num_layers, 4 * self.hidden_size, 8 * self.hidden_size, use_pool=False
+ )
+ blocks["output"] = nn.Sequential(output_blocks)
+
+ self.blocks = nn.Sequential(blocks)
+
+ self.post_init()
+
+ if self.config.freeze:
+ for param in self.parameters():
+ param.requires_grad = False
+
+ def get_codebook_indices(self, pixel_values: torch.Tensor) -> torch.Tensor:
+ f"""
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Pixel values. Codebook pixel values can be obtained using [`AutoImageProcessor`] by passing
+ `return_codebook_pixels=True`. See [`FlavaImageProcessor.__call__`] for details.
+
+ Examples:
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoImageProcessor, FlavaImageCodebook
+
+ >>> model = FlavaImageCodebook.from_pretrained("{_CHECKPOINT_FOR_CODEBOOK_DOC}")
+ >>> image_processor = AutoImageProcessor.from_pretrained("{_CHECKPOINT_FOR_CODEBOOK_DOC}")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> inputs = image_processor([image], return_codebook_pixels=True, return_tensors="pt")
+ >>> inputs = dict(pixel_values=inputs.codebook_pixel_values)
+
+ >>> outputs = model.get_codebook_indices(**inputs)
+ ```
+ """
+ z_logits = self.blocks(pixel_values)
+ return torch.argmax(z_logits, axis=1)
+
+ def get_codebook_probs(self, pixel_values: torch.Tensor) -> torch.Tensor:
+ z_logits = self.blocks(pixel_values)
+ return nn.Softmax(dim=1)(z_logits)
+
+ def forward(self, pixel_values: torch.FloatTensor, **kwargs) -> torch.Tensor:
+ f"""
+ Args:
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
+ Pixel values. Codebook pixel values can be obtained using [`AutoImageProcessor`] by passing
+ `return_codebook_pixels=True`. See [`FlavaImageProcessor.__call__`] for details.
+
+ Examples:
+
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoImageProcessor, FlavaImageCodebook
+
+ >>> model = FlavaImageCodebook.from_pretrained("{_CHECKPOINT_FOR_CODEBOOK_DOC}")
+ >>> image_processor = AutoImageProcessor.from_pretrained("{_CHECKPOINT_FOR_CODEBOOK_DOC}")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> inputs = image_processor([image], return_codebook_pixels=True, return_tensors="pt")
+ >>> inputs = dict(pixel_values=inputs.codebook_pixel_values)
+
+ >>> outputs = model(**inputs)
+ >>> print(outputs.shape)
+ (1, 196)
+ ```
+ """
+ if len(pixel_values.shape) != 4:
+ raise ValueError(f"input shape {pixel_values.shape} is not 4d")
+ if pixel_values.shape[1] != self.input_channels:
+ raise ValueError(f"input has {pixel_values.shape[1]} channels but model built for {self.input_channels}")
+ return self.blocks(pixel_values)
+
+
+class FlavaPredictionHeadTransform(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ if isinstance(config.hidden_act, str):
+ self.transform_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.transform_act_fn = config.hidden_act
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+ def forward(self, hidden_states):
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.transform_act_fn(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states)
+ return hidden_states
+
+
+class FlavaMaskedPredictionHead(nn.Module):
+ def __init__(self, config, weight=None):
+ super().__init__()
+ self.config = config
+ self.transform = FlavaPredictionHeadTransform(config)
+ self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=True)
+ self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+ if weight is not None:
+ self.decoder.weight = weight
+
+ def forward(self, x):
+ x = self.transform(x)
+ x = self.decoder(x)
+ return x
+
+
+class FlavaITMHead(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.pooler = FlavaPooler(config)
+ self.seq_relationship = nn.Linear(config.hidden_size, 2)
+
+ def forward(self, x):
+ x = self.pooler(x)
+ x = self.seq_relationship(x)
+ return x
+
+
+class FlavaGlobalContrastiveHead(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.global_backprop_contrastive = config.global_backprop_contrastive
+
+ def forward(self, image_embeddings, text_embeddings, logit_scale):
+ temperature = torch.exp(logit_scale)
+ if not torch.distributed.is_available() or not torch.distributed.is_initialized():
+ labels = torch.arange(image_embeddings.size(0), device=image_embeddings.device)
+ image_embeddings_all = [image_embeddings]
+ text_embeddings_all = [text_embeddings]
+ else:
+ local_batch_size = image_embeddings.size(0)
+ world_size = torch.distributed.get_world_size()
+
+ if self.global_backprop_contrastive:
+ # `torch.distributed.nn.functional.all_gather` does backprop on all active workers
+ # whereas `torch.distributed.all_gather` does only backpropagates on the current worker.
+ image_embeddings_all = torch.distributed.nn.functional.all_gather(image_embeddings)
+ text_embeddings_all = torch.distributed.nn.functional.all_gather(text_embeddings)
+ else:
+ image_embeddings_all = [torch.zeros_like(text_embeddings) for _ in range(world_size)]
+ text_embeddings_all = [torch.zeros_like(image_embeddings) for _ in range(world_size)]
+ torch.distributed.all_gather(image_embeddings_all, image_embeddings)
+ torch.distributed.all_gather(text_embeddings_all, text_embeddings)
+
+ labels = local_batch_size * torch.distributed.get_rank() + torch.arange(
+ local_batch_size, device=image_embeddings.device
+ )
+
+ image_embeddings_all = torch.cat(image_embeddings_all)
+ text_embeddings_all = torch.cat(text_embeddings_all)
+
+ logits_per_image = torch.matmul(image_embeddings, text_embeddings_all.transpose(0, 1)) * temperature
+ logits_per_text = torch.matmul(text_embeddings, image_embeddings_all.transpose(0, 1)) * temperature
+
+ return logits_per_image, logits_per_text, labels
+
+
+@auto_docstring(
+ custom_intro="""
+ The FLAVA model for pretraining which outputs losses, embeddings, logits and transformer outputs.
+ """
+)
+class FlavaForPreTraining(FlavaPreTrainedModel):
+ # Those are linked to xxx.bias
+ _tied_weights_keys = {
+ "mmm_text_head.bias": "mmm_text_head.decoder.bias",
+ "mim_head.bias": "mim_head.decoder.bias",
+ "mlm_head.bias": "mlm_head.decoder.bias",
+ "mmm_image_head.bias": "mmm_image_head.decoder.bias",
+ }
+
+ def __init__(self, config: FlavaConfig, image_codebook: nn.Module | None = None):
+ r"""
+ image_codebook ([`nn.Module`]):
+ If passed, the image codebook will be set to this. Otherwise, it will be initialized using the
+ image_codebook_config defined in the config first as the first parameter.
+ """
+ super().__init__(config)
+ self.flava = FlavaModel(config)
+
+ self.image_codebook = image_codebook
+ if self.image_codebook is None and config.init_codebook:
+ self.image_codebook = FlavaImageCodebook(config.image_codebook_config)
+
+ # Levarage text and image encoder configs to create the masked
+ # head since it has the right vocab
+ self.mim_head = FlavaMaskedPredictionHead(config.image_config)
+ self.mlm_head = FlavaMaskedPredictionHead(config.text_config)
+ self.itm_head = FlavaITMHead(config)
+ self.mmm_image_head = FlavaMaskedPredictionHead(config.image_config)
+ self.mmm_text_head = FlavaMaskedPredictionHead(config.text_config)
+ self.global_contrastive_head = FlavaGlobalContrastiveHead(config)
+
+ self.image_vocab_size = config.image_config.vocab_size
+ self.text_vocab_size = config.text_config.vocab_size
+ self.mlm_weight = config.mlm_weight
+ self.mim_weight = config.mim_weight
+ self.global_contrastive_weight = config.global_contrastive_weight
+ self.ce_ignore_index = config.ce_ignore_index
+ self.itm_weight = config.itm_weight
+ self.mmm_image_weight = config.mmm_image_weight
+ self.mmm_text_weight = config.mmm_text_weight
+ self.skip_unmasked_multimodal_encoder = config.skip_unmasked_multimodal_encoder
+
+ self.post_init()
+
+ def _resize_to_2d(self, x: torch.Tensor):
+ if x.dim() > 2:
+ x = x.view(x.size(0), -1)
+ return x
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ input_ids_masked: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ codebook_pixel_values: torch.FloatTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ bool_masked_pos: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ image_attention_mask: torch.Tensor | None = None,
+ skip_unmasked_multimodal_encoder: bool | None = None,
+ mlm_labels: torch.Tensor | None = None,
+ mim_labels: torch.Tensor | None = None,
+ itm_labels: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool = True,
+ return_dict: bool | None = None,
+ return_loss: bool | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor] | FlavaForPreTrainingOutput:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, text_seq_len)`):
+ Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`AutoTokenizer`]. See
+ [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input
+ IDs?](../glossary#input-ids)
+ input_ids_masked (`torch.LongTensor` of shape `(batch_size, text_seq_len)`):
+ Indices of input sequence tokens in the vocabulary. These ones are the masked version of the original task
+ to be used with MLM. Indices can be obtained using [`AutoTokenizer`] along with
+ [`DataCollatorForMaskedLanguageModeling`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids)
+ codebook_pixel_values (`torch.FloatTensor` of shape `(batch_size, num_image_patches, patch_size, patch_size, 3)`, *optional*):
+ Pixel values for image patches that are used to compute the image codebook labels for masked image modeling.
+ token_type_ids (`torch.LongTensor` of shape `(batch_size, text_seq_len)`, *optional*):
+ Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
+ 1]`:
+ - 0 corresponds to a *sentence A* token,
+ - 1 corresponds to a *sentence B* token.
+ [What are token type IDs?](../glossary#token-type-ids)
+ bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, image_num_patches)`):
+ Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
+ image_attention_mask (`torch.FloatTensor` of shape `(batch_size, image_num_patches)`, *optional*):
+ Mask to avoid performing attention on padding token indices specifically for images. 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)
+ skip_unmasked_multimodal_encoder (*bool*, *optional*):
+ Skip any calculations for multimodal encoder for unmasked inputs. FLAVA pretraining doesn't need unmasked
+ multimodal embeddings or outputs as of now.
+ mlm_labels (`torch.LongTensor` of shape `(batch_size, text_seq_len)`, *optional*):
+ Labels for computing the left-to-right language and multimodal masked modeling loss (next word prediction).
+ Indices should be in `[-100, 0, ..., text_config.vocab_size - 1]` (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,
+ ..., text_config.vocab_size - 1]`.
+ mim_labels (`torch.LongTensor` of shape `(batch_size, image_num_patches)`, *optional*):
+ Labels for computing the image and multimodal masked modeling loss. Indices should be in `[-100, 0, ...,
+ image_config.vocab_size - 1]`. Tokens with indices set to `-100` are ignored (masked), the loss is only
+ computed for the tokens with labels in `[0, ..., image_config.vocab_size - 1]`. If not passed, they are
+ generated automatically using the image codebook assigned to the model. By default, it uses
+ [`FlavaImageCodebook`]. See [`FlavaImageCodebook`] to understand how to generate mim_labels.
+ itm_labels (`torch.LongTensor` of shape `(batch_size, 1)`, *optional*):
+ Labels for computing the image-text matching loss. 0 means the pairs don't match and 1 means they match.
+ The pairs with 0 will be skipped for calculation of MMM and global contrastive losses as well.
+ return_loss (`bool`, *optional*, default to None):
+ Whether to return calculated loss or not.
+
+ Examples:
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import FlavaForPreTraining, AutoProcessor
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> model = FlavaForPreTraining.from_pretrained("facebook/flava-full")
+ >>> processor = AutoProcessor.from_pretrained("facebook/flava-full")
+
+ >>> text = ["a photo of a cat"]
+
+ >>> inputs = processor(
+ ... images=[image],
+ ... text=text,
+ ... return_masks=True,
+ ... return_codebook_pixels=True,
+ ... padding=True,
+ ... max_length=77,
+ ... return_tensors="pt",
+ ... )
+
+
+ >>> output = model(**inputs)
+ ```
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+ return_loss = return_loss if return_loss is not None else self.config.return_loss
+
+ skip_unmasked_multimodal_encoder = (
+ skip_unmasked_multimodal_encoder
+ if skip_unmasked_multimodal_encoder is not None
+ else self.skip_unmasked_multimodal_encoder
+ )
+
+ if input_ids_masked is None and input_ids is not None:
+ logger.warning(
+ "`input_ids_masked` isn't passed which means MLM loss won't be calculated correctlySetting it to"
+ " `input_ids` so that model can work. Please pass it if this is unintentional. This is usually OKAY if"
+ " you are doing inference on unmasked text..."
+ )
+ input_ids_masked = input_ids
+
+ flava_output = self.flava(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ image_attention_mask=image_attention_mask,
+ # Don't need unmasked multimodal embedding for anything so skip it
+ # NOTE: ITM uses masked version
+ skip_multimodal_encoder=skip_unmasked_multimodal_encoder,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ # Pass true to have deterministic outputs
+ return_dict=True,
+ )
+
+ flava_masked_output = self.flava(
+ input_ids=input_ids_masked,
+ pixel_values=pixel_values,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ image_attention_mask=image_attention_mask,
+ bool_masked_pos=bool_masked_pos,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=True,
+ )
+
+ pos_mask = None
+
+ image_embeddings = flava_output.image_embeddings
+ text_embeddings = flava_output.text_embeddings
+ image_masked_embeddings = flava_masked_output.image_embeddings
+ text_masked_embeddings = flava_masked_output.text_embeddings
+ multimodal_masked_embeddings = flava_masked_output.multimodal_embeddings
+
+ total_loss = mim_loss = mlm_loss = mmm_text_loss = mmm_image_loss = gc_loss = itm_loss = None
+ mim_logits = mlm_logits = mmm_text_logits = mmm_image_logits = None
+ itm_logits = logits_per_image = logits_per_text = None
+
+ # Calculate mim_labels if necessary from the image_codebook
+ if image_masked_embeddings is not None or multimodal_masked_embeddings is not None:
+ if mim_labels is None and return_loss:
+ if self.image_codebook is None:
+ raise RuntimeError(
+ "`return_loss` is set to True but the image codebook is not initialized and no `mim_labels` "
+ " have been passed. Reinstantiate the model with `init_codebook` set to True or "
+ "pass in your custom `mim_labels`"
+ )
+ if codebook_pixel_values is None:
+ raise ValueError(
+ "`codebook_pixel_value` are required to generate `mim_labels` if loss is expected. "
+ "Call `AutoProcessor` with `return_codebook_pixels` set to True"
+ )
+ mim_labels = self.image_codebook.get_codebook_indices(codebook_pixel_values)
+ # Unimodal MIM Loss
+ # If multimodal embeddings are present, we will calculate MMM loss
+ if self.mim_weight > 0 and image_masked_embeddings is not None and multimodal_masked_embeddings is None:
+ sequence_for_image = image_masked_embeddings
+
+ if mim_labels is not None:
+ mim_labels = self._resize_to_2d(mim_labels)
+ bool_masked_pos = self._resize_to_2d(bool_masked_pos)
+ mim_labels[bool_masked_pos.ne(True)] = self.ce_ignore_index
+
+ sequence_for_image = sequence_for_image[:, -mim_labels.size(1) :, :]
+ masked_tokens = mim_labels.ne(self.ce_ignore_index)
+ mim_labels_filtered = mim_labels[masked_tokens]
+ sequence_for_image = sequence_for_image[masked_tokens, :]
+ mim_logits = self.mim_head(sequence_for_image)
+ if return_loss:
+ mim_loss = nn.functional.cross_entropy(
+ mim_logits.view(-1, self.image_vocab_size), mim_labels_filtered.view(-1)
+ )
+ mim_loss *= self.mim_weight
+ else:
+ mim_logits = self.mim_head(sequence_for_image)
+
+ # Unimodal MLM Loss
+ if self.mlm_weight > 0 and text_masked_embeddings is not None and multimodal_masked_embeddings is None:
+ sequence_for_text = text_masked_embeddings
+ if mlm_labels is not None:
+ mlm_labels = self._resize_to_2d(mlm_labels)
+ sequence_for_text = sequence_for_text[:, -mlm_labels.size(1) :, :]
+ masked_tokens = mlm_labels.ne(self.ce_ignore_index)
+ mlm_labels_filtered = mlm_labels[masked_tokens]
+ sequence_for_text = sequence_for_text[masked_tokens, :]
+ mlm_logits = self.mlm_head(sequence_for_text)
+ if return_loss:
+ mlm_loss = nn.functional.cross_entropy(
+ mlm_logits.view(-1, self.text_vocab_size), mlm_labels_filtered.view(-1)
+ )
+ mlm_loss *= self.mlm_weight
+ else:
+ mlm_logits = self.mlm_head(sequence_for_text)
+
+ # ITM Loss
+ if self.itm_weight > 0 and multimodal_masked_embeddings is not None:
+ itm_logits = self.itm_head(multimodal_masked_embeddings)
+
+ if itm_labels is not None:
+ pos_pairs = itm_labels.ne(0)
+ pos_mask = torch.where(pos_pairs.any(), pos_pairs, pos_pairs.new([True]))
+ if return_loss:
+ itm_loss = nn.functional.cross_entropy(itm_logits, itm_labels)
+ itm_loss *= self.itm_weight
+
+ if multimodal_masked_embeddings is not None:
+ multimodal_masked_embeddings = multimodal_masked_embeddings[pos_mask]
+
+ if mlm_labels is not None:
+ mlm_labels = mlm_labels[pos_mask]
+
+ if mim_labels is not None:
+ mim_labels = mim_labels[pos_mask]
+ bool_masked_pos = bool_masked_pos[pos_mask]
+
+ # MMM Image Loss
+ if multimodal_masked_embeddings is not None and self.mmm_image_weight > 0:
+ sequence_for_image = multimodal_masked_embeddings
+ end_index = image_masked_embeddings.size(1) - 1
+ sequence_for_image = sequence_for_image[:, 2 : 2 + end_index, :]
+
+ if mim_labels is not None:
+ mim_labels = self._resize_to_2d(mim_labels)
+ bool_masked_pos = self._resize_to_2d(bool_masked_pos)
+ mim_labels[bool_masked_pos.ne(True)] = self.ce_ignore_index
+
+ masked_tokens = mim_labels.ne(self.ce_ignore_index)
+ mim_labels_filtered = mim_labels[masked_tokens]
+ sequence_for_image = sequence_for_image[masked_tokens, :]
+ mmm_image_logits = self.mmm_image_head(sequence_for_image)
+ if return_loss:
+ mmm_image_loss = nn.functional.cross_entropy(
+ mmm_image_logits.view(-1, self.image_vocab_size), mim_labels_filtered.view(-1)
+ )
+ mmm_image_loss *= self.mmm_image_weight
+ else:
+ mmm_image_logits = self.mmm_image_head(sequence_for_image)
+
+ # MMM Text Loss
+ if multimodal_masked_embeddings is not None and self.mmm_text_weight > 0:
+ sequence_for_text = multimodal_masked_embeddings
+ sequence_for_text = sequence_for_text[:, -text_masked_embeddings.size(1) :, :]
+
+ if mlm_labels is not None:
+ mlm_labels = self._resize_to_2d(mlm_labels)
+ masked_tokens = mlm_labels.ne(self.ce_ignore_index)
+ mlm_labels_filtered = mlm_labels[masked_tokens]
+ sequence_for_text = sequence_for_text[masked_tokens, :]
+ mmm_text_logits = self.mmm_text_head(sequence_for_text)
+ if return_loss:
+ mmm_text_loss = nn.functional.cross_entropy(
+ mmm_text_logits.view(-1, self.text_vocab_size), mlm_labels_filtered.view(-1)
+ )
+ mmm_text_loss *= self.mmm_text_weight
+ else:
+ mmm_text_logits = self.mmm_text_head(sequence_for_text)
+
+ # Global Contrastive Loss
+ if image_embeddings is not None and text_embeddings is not None and self.global_contrastive_weight > 0:
+ text_embedding = self.flava.text_projection(text_embeddings[:, 0, :])
+ text_embedding = nn.functional.normalize(text_embedding, dim=-1)
+
+ image_embedding = self.flava.image_projection(image_embeddings[:, 0, :])
+ image_embedding = nn.functional.normalize(image_embedding, dim=-1)
+
+ if self.training:
+ self.flava.logit_scale.data.clamp_(LOGIT_SCALE_CLAMP_MIN, LOGIT_SCALE_CLAMP_MAX)
+
+ logits_per_image, logits_per_text, gc_labels = self.global_contrastive_head(
+ image_embedding, text_embedding, self.flava.logit_scale
+ )
+
+ # Apply ITM negative mask if any
+ if pos_mask is not None:
+ logits_per_image = logits_per_image[pos_mask]
+ logits_per_text = logits_per_text[pos_mask]
+ gc_labels = gc_labels[pos_mask]
+
+ if return_loss:
+ gc_loss_image = nn.functional.cross_entropy(logits_per_image, gc_labels)
+ gc_loss_text = nn.functional.cross_entropy(logits_per_text, gc_labels)
+ gc_loss = (gc_loss_image + gc_loss_text) / 2
+ gc_loss *= self.global_contrastive_weight
+
+ flava_losses = FlavaLosses(
+ mim=mim_loss,
+ mlm=mlm_loss,
+ itm=itm_loss,
+ global_contrastive=gc_loss,
+ mmm_image=mmm_image_loss,
+ mmm_text=mmm_text_loss,
+ )
+
+ if return_loss and not flava_losses.all_none():
+ total_loss = sum(loss if loss is not None else 0 for loss in flava_losses.values())
+
+ if not return_dict:
+ output = (
+ image_embeddings,
+ flava_output.image_output.to_tuple() if flava_output.image_output is not None else None,
+ text_embeddings,
+ flava_output.text_output.to_tuple() if flava_output.text_output is not None else None,
+ flava_output.multimodal_embeddings,
+ flava_output.multimodal_output.to_tuple() if flava_output.multimodal_output is not None else None,
+ image_masked_embeddings,
+ flava_masked_output.image_output.to_tuple() if flava_masked_output.image_output is not None else None,
+ text_masked_embeddings,
+ flava_masked_output.text_output.to_tuple() if flava_masked_output.text_output is not None else None,
+ multimodal_masked_embeddings,
+ flava_masked_output.multimodal_output.to_tuple()
+ if flava_masked_output.multimodal_output is not None
+ else None,
+ mim_logits,
+ mlm_logits,
+ itm_logits,
+ logits_per_image,
+ logits_per_image,
+ mmm_image_logits,
+ mmm_text_logits,
+ )
+ if return_loss and not flava_losses.all_none():
+ output = (
+ total_loss,
+ flava_losses,
+ ) + output
+
+ # Filter None as transformer by default won't handle it
+ return tuple(x for x in output if x is None)
+
+ return FlavaForPreTrainingOutput(
+ loss=total_loss,
+ loss_info=flava_losses,
+ image_embeddings=image_embeddings,
+ image_output=flava_output.image_output,
+ text_embeddings=text_embeddings,
+ text_output=flava_output.text_output,
+ multimodal_embeddings=flava_output.multimodal_embeddings,
+ multimodal_output=flava_output.multimodal_output,
+ image_masked_embeddings=image_masked_embeddings,
+ image_masked_output=flava_masked_output.image_output,
+ text_masked_embeddings=text_masked_embeddings,
+ text_masked_output=flava_masked_output.text_output,
+ multimodal_masked_embeddings=multimodal_masked_embeddings,
+ multimodal_masked_output=flava_masked_output.multimodal_output,
+ mim_logits=mim_logits,
+ mlm_logits=mlm_logits,
+ itm_logits=itm_logits,
+ contrastive_logits_per_image=logits_per_image,
+ contrastive_logits_per_text=logits_per_text,
+ mmm_image_logits=mmm_image_logits,
+ mmm_text_logits=mmm_text_logits,
+ )
+
+
+__all__ = [
+ "FlavaForPreTraining",
+ "FlavaImageCodebook",
+ "FlavaImageModel",
+ "FlavaModel",
+ "FlavaMultimodalModel",
+ "FlavaPreTrainedModel",
+ "FlavaTextModel",
+]
diff --git a/third_party/transformers/src/transformers/models/flava/processing_flava.py b/third_party/transformers/src/transformers/models/flava/processing_flava.py
new file mode 100644
index 0000000000000000000000000000000000000000..592654a18f0e9fc07fe3c1a1e1ec8abb3459c3ff
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/flava/processing_flava.py
@@ -0,0 +1,28 @@
+# Copyright 2022 Meta Platforms authors and The HuggingFace Team. All rights reserved.
+#
+# 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.
+"""
+Image/Text processor class for FLAVA
+"""
+
+from ...processing_utils import ProcessorMixin
+from ...utils import auto_docstring
+
+
+@auto_docstring
+class FlavaProcessor(ProcessorMixin):
+ def __init__(self, image_processor=None, tokenizer=None, **kwargs):
+ super().__init__(image_processor, tokenizer)
+
+
+__all__ = ["FlavaProcessor"]
diff --git a/third_party/transformers/src/transformers/models/fnet/__init__.py b/third_party/transformers/src/transformers/models/fnet/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..00785b5e50753c20f5049ce1f74a5d4878ab7499
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/fnet/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_fnet import *
+ from .modeling_fnet import *
+ from .tokenization_fnet import FNetTokenizer, FNetTokenizerFast
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/fnet/configuration_fnet.py b/third_party/transformers/src/transformers/models/fnet/configuration_fnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..08fb9a88910551aa3ddb155690ad63ee834d539c
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/fnet/configuration_fnet.py
@@ -0,0 +1,69 @@
+# Copyright 2021 Google AI and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""FNet model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="google/fnet-base")
+@strict
+class FNetConfig(PreTrainedConfig):
+ r"""
+ use_tpu_fourier_optimizations (`bool`, *optional*, defaults to `False`):
+ Determines whether to use TPU optimized FFTs. If `True`, the model will favor axis-wise FFTs transforms.
+ Set to `False` for GPU/CPU hardware, in which case n-dimensional FFTs are used.
+ tpu_short_seq_length (`int`, *optional*, defaults to 512):
+ The sequence length that is expected by the model when using TPUs. This will be used to initialize the DFT
+ matrix only when *use_tpu_fourier_optimizations* is set to `True` and the input sequence is shorter than or
+ equal to 4096 tokens.
+
+ Example:
+
+ ```python
+ >>> from transformers import FNetConfig, FNetModel
+
+ >>> # Initializing a FNet fnet-base style configuration
+ >>> configuration = FNetConfig()
+
+ >>> # Initializing a model (with random weights) from the fnet-base style configuration
+ >>> model = FNetModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "fnet"
+
+ vocab_size: int = 32000
+ hidden_size: int = 768
+ num_hidden_layers: int = 12
+ intermediate_size: int = 3072
+ hidden_act: str = "gelu_new"
+ hidden_dropout_prob: float | int = 0.1
+ max_position_embeddings: int = 512
+ type_vocab_size: int = 4
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-12
+ use_tpu_fourier_optimizations: bool = False
+ tpu_short_seq_length: int = 512
+ pad_token_id: int | None = 3
+ bos_token_id: int | None = 1
+ eos_token_id: int | list[int] | None = 2
+ tie_word_embeddings: bool = True
+
+
+__all__ = ["FNetConfig"]
diff --git a/third_party/transformers/src/transformers/models/fnet/convert_fnet_original_flax_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/fnet/convert_fnet_original_flax_checkpoint_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..3f94e360f1877ba2ed818593a6196630c228505e
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/fnet/convert_fnet_original_flax_checkpoint_to_pytorch.py
@@ -0,0 +1,155 @@
+# Copyright 2021 The HuggingFace Inc. team.
+#
+# 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.
+"""Convert FNet checkpoint."""
+
+import argparse
+
+import torch
+from flax.training.checkpoints import restore_checkpoint
+
+from transformers import FNetConfig, FNetForPreTraining
+from transformers.utils import logging
+
+
+logging.set_verbosity_info()
+
+
+def convert_flax_checkpoint_to_pytorch(flax_checkpoint_path, fnet_config_file, save_path):
+ # Initialise PyTorch model
+ config = FNetConfig.from_json_file(fnet_config_file)
+ print(f"Building PyTorch model from configuration: {config}")
+ fnet_pretraining_model = FNetForPreTraining(config)
+
+ checkpoint_dict = restore_checkpoint(flax_checkpoint_path, None)
+ pretrained_model_params = checkpoint_dict["target"]
+
+ # Embeddings
+ # Position IDs
+ state_dict = fnet_pretraining_model.state_dict()
+
+ position_ids = state_dict["fnet.embeddings.position_ids"]
+ new_state_dict = {"fnet.embeddings.position_ids": position_ids}
+ # Embedding Layers
+ new_state_dict["fnet.embeddings.word_embeddings.weight"] = torch.tensor(
+ pretrained_model_params["encoder"]["embedder"]["word"]["embedding"]
+ )
+ new_state_dict["fnet.embeddings.position_embeddings.weight"] = torch.tensor(
+ pretrained_model_params["encoder"]["embedder"]["position"]["embedding"][0]
+ )
+ new_state_dict["fnet.embeddings.token_type_embeddings.weight"] = torch.tensor(
+ pretrained_model_params["encoder"]["embedder"]["type"]["embedding"]
+ )
+ new_state_dict["fnet.embeddings.projection.weight"] = torch.tensor(
+ pretrained_model_params["encoder"]["embedder"]["hidden_mapping_in"]["kernel"]
+ ).T
+ new_state_dict["fnet.embeddings.projection.bias"] = torch.tensor(
+ pretrained_model_params["encoder"]["embedder"]["hidden_mapping_in"]["bias"]
+ )
+ new_state_dict["fnet.embeddings.LayerNorm.weight"] = torch.tensor(
+ pretrained_model_params["encoder"]["embedder"]["layer_norm"]["scale"]
+ )
+ new_state_dict["fnet.embeddings.LayerNorm.bias"] = torch.tensor(
+ pretrained_model_params["encoder"]["embedder"]["layer_norm"]["bias"]
+ )
+
+ # Encoder Layers
+ for layer in range(config.num_hidden_layers):
+ new_state_dict[f"fnet.encoder.layer.{layer}.fourier.output.LayerNorm.weight"] = torch.tensor(
+ pretrained_model_params["encoder"][f"encoder_{layer}"]["mixing_layer_norm"]["scale"]
+ )
+ new_state_dict[f"fnet.encoder.layer.{layer}.fourier.output.LayerNorm.bias"] = torch.tensor(
+ pretrained_model_params["encoder"][f"encoder_{layer}"]["mixing_layer_norm"]["bias"]
+ )
+
+ new_state_dict[f"fnet.encoder.layer.{layer}.intermediate.dense.weight"] = torch.tensor(
+ pretrained_model_params["encoder"][f"feed_forward_{layer}"]["intermediate"]["kernel"]
+ ).T
+ new_state_dict[f"fnet.encoder.layer.{layer}.intermediate.dense.bias"] = torch.tensor(
+ pretrained_model_params["encoder"][f"feed_forward_{layer}"]["intermediate"]["bias"]
+ )
+
+ new_state_dict[f"fnet.encoder.layer.{layer}.output.dense.weight"] = torch.tensor(
+ pretrained_model_params["encoder"][f"feed_forward_{layer}"]["output"]["kernel"]
+ ).T
+ new_state_dict[f"fnet.encoder.layer.{layer}.output.dense.bias"] = torch.tensor(
+ pretrained_model_params["encoder"][f"feed_forward_{layer}"]["output"]["bias"]
+ )
+
+ new_state_dict[f"fnet.encoder.layer.{layer}.output.LayerNorm.weight"] = torch.tensor(
+ pretrained_model_params["encoder"][f"encoder_{layer}"]["output_layer_norm"]["scale"]
+ )
+ new_state_dict[f"fnet.encoder.layer.{layer}.output.LayerNorm.bias"] = torch.tensor(
+ pretrained_model_params["encoder"][f"encoder_{layer}"]["output_layer_norm"]["bias"]
+ )
+
+ # Pooler Layers
+ new_state_dict["fnet.pooler.dense.weight"] = torch.tensor(pretrained_model_params["encoder"]["pooler"]["kernel"]).T
+ new_state_dict["fnet.pooler.dense.bias"] = torch.tensor(pretrained_model_params["encoder"]["pooler"]["bias"])
+
+ # Masked LM Layers
+ new_state_dict["cls.predictions.transform.dense.weight"] = torch.tensor(
+ pretrained_model_params["predictions_dense"]["kernel"]
+ ).T
+ new_state_dict["cls.predictions.transform.dense.bias"] = torch.tensor(
+ pretrained_model_params["predictions_dense"]["bias"]
+ )
+ new_state_dict["cls.predictions.transform.LayerNorm.weight"] = torch.tensor(
+ pretrained_model_params["predictions_layer_norm"]["scale"]
+ )
+ new_state_dict["cls.predictions.transform.LayerNorm.bias"] = torch.tensor(
+ pretrained_model_params["predictions_layer_norm"]["bias"]
+ )
+ new_state_dict["cls.predictions.decoder.weight"] = torch.tensor(
+ pretrained_model_params["encoder"]["embedder"]["word"]["embedding"]
+ )
+ new_state_dict["cls.predictions.decoder.bias"] = torch.tensor(
+ pretrained_model_params["predictions_output"]["output_bias"]
+ )
+ new_state_dict["cls.predictions.bias"] = torch.tensor(pretrained_model_params["predictions_output"]["output_bias"])
+
+ # Seq Relationship Layers
+ new_state_dict["cls.seq_relationship.weight"] = torch.tensor(
+ pretrained_model_params["classification"]["output_kernel"]
+ )
+ new_state_dict["cls.seq_relationship.bias"] = torch.tensor(
+ pretrained_model_params["classification"]["output_bias"]
+ )
+
+ # Load State Dict
+ fnet_pretraining_model.load_state_dict(new_state_dict)
+
+ # Save PreTrained
+ print(f"Saving pretrained model to {save_path}")
+ fnet_pretraining_model.save_pretrained(save_path)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ # Required parameters
+ parser.add_argument(
+ "--flax_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
+ )
+ parser.add_argument(
+ "--fnet_config_file",
+ default=None,
+ type=str,
+ required=True,
+ help=(
+ "The config json file corresponding to the pre-trained FNet model. \n"
+ "This specifies the model architecture."
+ ),
+ )
+ parser.add_argument("--save_path", default=None, type=str, required=True, help="Path to the output model.")
+ args = parser.parse_args()
+ convert_flax_checkpoint_to_pytorch(args.flax_checkpoint_path, args.fnet_config_file, args.save_path)
diff --git a/third_party/transformers/src/transformers/models/fnet/modeling_fnet.py b/third_party/transformers/src/transformers/models/fnet/modeling_fnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..c5b08bd94050e82247bd13286832324e724ddf21
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/fnet/modeling_fnet.py
@@ -0,0 +1,1070 @@
+# Copyright 2021 Google Research and The HuggingFace Inc. team. All rights reserved.
+#
+# 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 FNet model."""
+
+from dataclasses import dataclass
+from functools import partial
+
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ... import initialization as init
+from ...utils import auto_docstring, is_scipy_available
+
+
+if is_scipy_available():
+ from scipy import linalg
+
+from ...activations import ACT2FN
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPooling,
+ MaskedLMOutput,
+ ModelOutput,
+ MultipleChoiceModelOutput,
+ NextSentencePredictorOutput,
+ QuestionAnsweringModelOutput,
+ SequenceClassifierOutput,
+ TokenClassifierOutput,
+)
+from ...modeling_utils import PreTrainedModel
+from ...pytorch_utils import apply_chunking_to_forward
+from ...utils import logging
+from .configuration_fnet import FNetConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# Adapted from https://github.com/google-research/google-research/blob/master/f_net/fourier.py
+def _two_dim_matmul(x, matrix_dim_one, matrix_dim_two):
+ """Applies 2D matrix multiplication to 3D input arrays."""
+ seq_length = x.shape[1]
+ matrix_dim_one = matrix_dim_one[:seq_length, :seq_length]
+ x = x.type(torch.complex64)
+ return torch.einsum("bij,jk,ni->bnk", x, matrix_dim_two, matrix_dim_one)
+
+
+# # Adapted from https://github.com/google-research/google-research/blob/master/f_net/fourier.py
+def two_dim_matmul(x, matrix_dim_one, matrix_dim_two):
+ return _two_dim_matmul(x, matrix_dim_one, matrix_dim_two)
+
+
+# Adapted from https://github.com/google-research/google-research/blob/master/f_net/fourier.py
+def fftn(x):
+ """
+ Applies n-dimensional Fast Fourier Transform (FFT) to input array.
+
+ Args:
+ x: Input n-dimensional array.
+
+ Returns:
+ n-dimensional Fourier transform of input n-dimensional array.
+ """
+ out = x
+ for axis in reversed(range(x.ndim)[1:]): # We don't need to apply FFT to last axis
+ out = torch.fft.fft(out, axis=axis)
+ return out
+
+
+class FNetEmbeddings(nn.Module):
+ """Construct the embeddings from word, position and token_type embeddings."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+ self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ # NOTE: This is the project layer and will be needed. The original code allows for different embedding and different model dimensions.
+ self.projection = nn.Linear(config.hidden_size, config.hidden_size)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+ self.register_buffer(
+ "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+ )
+
+ self.register_buffer(
+ "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
+ )
+
+ def forward(self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None):
+ if input_ids is not None:
+ input_shape = input_ids.size()
+ else:
+ input_shape = inputs_embeds.size()[:-1]
+
+ seq_length = input_shape[1]
+
+ if position_ids is None:
+ position_ids = self.position_ids[:, :seq_length]
+
+ # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
+ # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
+ # issue #5664
+ if token_type_ids is None:
+ if hasattr(self, "token_type_ids"):
+ buffered_token_type_ids = self.token_type_ids[:, :seq_length]
+ buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
+ token_type_ids = buffered_token_type_ids_expanded
+ else:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.word_embeddings(input_ids)
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+ embeddings = inputs_embeds + token_type_embeddings
+
+ position_embeddings = self.position_embeddings(position_ids)
+ embeddings += position_embeddings
+ embeddings = self.LayerNorm(embeddings)
+ embeddings = self.projection(embeddings)
+ embeddings = self.dropout(embeddings)
+ return embeddings
+
+
+class FNetBasicFourierTransform(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self._init_fourier_transform(config)
+
+ def _init_fourier_transform(self, config):
+ if not config.use_tpu_fourier_optimizations:
+ self.fourier_transform = partial(torch.fft.fftn, dim=(1, 2))
+ elif config.max_position_embeddings <= 4096:
+ if is_scipy_available():
+ self.register_buffer(
+ "dft_mat_hidden", torch.tensor(linalg.dft(config.hidden_size), dtype=torch.complex64)
+ )
+ self.register_buffer(
+ "dft_mat_seq", torch.tensor(linalg.dft(config.tpu_short_seq_length), dtype=torch.complex64)
+ )
+ self.fourier_transform = partial(
+ two_dim_matmul, matrix_dim_one=self.dft_mat_seq, matrix_dim_two=self.dft_mat_hidden
+ )
+ else:
+ logging.warning(
+ "SciPy is needed for DFT matrix calculation and is not found. Using TPU optimized fast fourier"
+ " transform instead."
+ )
+ self.fourier_transform = fftn
+ else:
+ self.fourier_transform = fftn
+
+ def forward(self, hidden_states):
+ # NOTE: We do not use torch.vmap as it is not integrated into PyTorch stable versions.
+ # Interested users can modify the code to use vmap from the nightly versions, getting the vmap from here:
+ # https://pytorch.org/docs/master/generated/torch.vmap.html. Note that fourier transform methods will need
+ # change accordingly.
+
+ outputs = self.fourier_transform(hidden_states).real
+ return (outputs,)
+
+
+class FNetBasicOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+ def forward(self, hidden_states, input_tensor):
+ hidden_states = self.LayerNorm(input_tensor + hidden_states)
+ return hidden_states
+
+
+class FNetFourierTransform(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.self = FNetBasicFourierTransform(config)
+ self.output = FNetBasicOutput(config)
+
+ def forward(self, hidden_states):
+ self_outputs = self.self(hidden_states)
+ fourier_output = self.output(self_outputs[0], hidden_states)
+ outputs = (fourier_output,)
+ return outputs
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->FNet
+class FNetIntermediate(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->FNet
+class FNetOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+class FNetLayer(GradientCheckpointingLayer):
+ def __init__(self, config):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1 # The dimension which has the sequence length
+ self.fourier = FNetFourierTransform(config)
+ self.intermediate = FNetIntermediate(config)
+ self.output = FNetOutput(config)
+
+ def forward(self, hidden_states):
+ self_fourier_outputs = self.fourier(hidden_states)
+ fourier_output = self_fourier_outputs[0]
+
+ layer_output = apply_chunking_to_forward(
+ self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, fourier_output
+ )
+
+ outputs = (layer_output,)
+
+ return outputs
+
+ def feed_forward_chunk(self, fourier_output):
+ intermediate_output = self.intermediate(fourier_output)
+ layer_output = self.output(intermediate_output, fourier_output)
+ return layer_output
+
+
+class FNetEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList([FNetLayer(config) for _ in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+
+ def forward(self, hidden_states, output_hidden_states=False, return_dict=True):
+ all_hidden_states = () if output_hidden_states else None
+
+ for i, layer_module in enumerate(self.layer):
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ layer_outputs = layer_module(hidden_states)
+
+ hidden_states = layer_outputs[0]
+
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ if not return_dict:
+ return tuple(v for v in [hidden_states, all_hidden_states] if v is not None)
+
+ return BaseModelOutput(last_hidden_state=hidden_states, hidden_states=all_hidden_states)
+
+
+# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->FNet
+class FNetPooler(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.activation = nn.Tanh()
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ # We "pool" the model by simply taking the hidden state corresponding
+ # to the first token.
+ first_token_tensor = hidden_states[:, 0]
+ pooled_output = self.dense(first_token_tensor)
+ pooled_output = self.activation(pooled_output)
+ return pooled_output
+
+
+# Copied from transformers.models.bert.modeling_bert.BertPredictionHeadTransform with Bert->FNet
+class FNetPredictionHeadTransform(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ if isinstance(config.hidden_act, str):
+ self.transform_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.transform_act_fn = config.hidden_act
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.transform_act_fn(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states)
+ return hidden_states
+
+
+class FNetLMPredictionHead(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.transform = FNetPredictionHeadTransform(config)
+ self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
+ self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+ def forward(self, hidden_states):
+ hidden_states = self.transform(hidden_states)
+ hidden_states = self.decoder(hidden_states)
+ return hidden_states
+
+
+class FNetOnlyMLMHead(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.predictions = FNetLMPredictionHead(config)
+
+ def forward(self, sequence_output):
+ prediction_scores = self.predictions(sequence_output)
+ return prediction_scores
+
+
+# Copied from transformers.models.bert.modeling_bert.BertOnlyNSPHead with Bert->FNet
+class FNetOnlyNSPHead(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.seq_relationship = nn.Linear(config.hidden_size, 2)
+
+ def forward(self, pooled_output):
+ seq_relationship_score = self.seq_relationship(pooled_output)
+ return seq_relationship_score
+
+
+# Copied from transformers.models.bert.modeling_bert.BertPreTrainingHeads with Bert->FNet
+class FNetPreTrainingHeads(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.predictions = FNetLMPredictionHead(config)
+ self.seq_relationship = nn.Linear(config.hidden_size, 2)
+
+ def forward(self, sequence_output, pooled_output):
+ prediction_scores = self.predictions(sequence_output)
+ seq_relationship_score = self.seq_relationship(pooled_output)
+ return prediction_scores, seq_relationship_score
+
+
+@auto_docstring
+class FNetPreTrainedModel(PreTrainedModel):
+ config: FNetConfig
+ base_model_prefix = "fnet"
+ supports_gradient_checkpointing = True
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, FNetEmbeddings):
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+ init.zeros_(module.token_type_ids)
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Output type of [`FNetForPreTraining`].
+ """
+)
+class FNetForPreTrainingOutput(ModelOutput):
+ r"""
+ loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
+ Total loss as the sum of the masked language modeling loss and the next sequence prediction
+ (classification) loss.
+ 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).
+ seq_relationship_logits (`torch.FloatTensor` of shape `(batch_size, 2)`):
+ Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation
+ before SoftMax).
+ """
+
+ loss: torch.FloatTensor | None = None
+ prediction_logits: torch.FloatTensor | None = None
+ seq_relationship_logits: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring
+class FNetModel(FNetPreTrainedModel):
+ """
+
+ The model can behave as an encoder, following the architecture described in [FNet: Mixing Tokens with Fourier
+ Transforms](https://huggingface.co/papers/2105.03824) by James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, Santiago Ontanon.
+
+ """
+
+ def __init__(self, config, add_pooling_layer=True):
+ r"""
+ add_pooling_layer (bool, *optional*, defaults to `True`):
+ Whether to add a pooling layer
+ """
+ super().__init__(config)
+ self.config = config
+
+ self.embeddings = FNetEmbeddings(config)
+ self.encoder = FNetEncoder(config)
+
+ self.pooler = FNetPooler(config) if add_pooling_layer else None
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embeddings.word_embeddings
+
+ def set_input_embeddings(self, value):
+ self.embeddings.word_embeddings = value
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | BaseModelOutput:
+ 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.return_dict
+
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+ elif input_ids is not None:
+ input_shape = input_ids.size()
+ batch_size, seq_length = input_shape
+ elif inputs_embeds is not None:
+ input_shape = inputs_embeds.size()[:-1]
+ batch_size, seq_length = input_shape
+ else:
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+ if (
+ self.config.use_tpu_fourier_optimizations
+ and seq_length <= 4096
+ and self.config.tpu_short_seq_length != seq_length
+ ):
+ raise ValueError(
+ "The `tpu_short_seq_length` in FNetConfig should be set equal to the sequence length being passed to"
+ " the model when using TPU optimizations."
+ )
+
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+ if token_type_ids is None:
+ if hasattr(self.embeddings, "token_type_ids"):
+ buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
+ buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
+ token_type_ids = buffered_token_type_ids_expanded
+ else:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+ embedding_output = self.embeddings(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ )
+ encoder_outputs = self.encoder(
+ embedding_output,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ sequence_output = encoder_outputs[0]
+
+ pooler_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+ if not return_dict:
+ return (sequence_output, pooler_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=sequence_output,
+ pooler_output=pooler_output,
+ hidden_states=encoder_outputs.hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ FNet Model with two heads on top as done during the pretraining: a `masked language modeling` head and a `next
+ sentence prediction (classification)` head.
+ """
+)
+class FNetForPreTraining(FNetPreTrainedModel):
+ _tied_weights_keys = {
+ "cls.predictions.decoder.bias": "cls.predictions.bias",
+ "cls.predictions.decoder.weight": "fnet.embeddings.word_embeddings.weight",
+ }
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.fnet = FNetModel(config)
+ self.cls = FNetPreTrainingHeads(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.cls.predictions.decoder
+
+ def set_output_embeddings(self, new_embeddings):
+ self.cls.predictions.decoder = new_embeddings
+ self.cls.predictions.bias = new_embeddings.bias
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ next_sentence_label: torch.Tensor | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | FNetForPreTrainingOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
+ config.vocab_size]` (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]`
+ next_sentence_label (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair
+ (see `input_ids` docstring) Indices should be in `[0, 1]`:
+
+ - 0 indicates sequence B is a continuation of sequence A,
+ - 1 indicates sequence B is a random sequence.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, FNetForPreTraining
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("google/fnet-base")
+ >>> model = FNetForPreTraining.from_pretrained("google/fnet-base")
+ >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+ >>> outputs = model(**inputs)
+ >>> prediction_logits = outputs.prediction_logits
+ >>> seq_relationship_logits = outputs.seq_relationship_logits
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.fnet(
+ input_ids,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ sequence_output, pooled_output = outputs[:2]
+ prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)
+
+ total_loss = None
+ if labels is not None and next_sentence_label is not None:
+ loss_fct = CrossEntropyLoss()
+ masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+ next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
+ total_loss = masked_lm_loss + next_sentence_loss
+
+ if not return_dict:
+ output = (prediction_scores, seq_relationship_score) + outputs[2:]
+ return ((total_loss,) + output) if total_loss is not None else output
+
+ return FNetForPreTrainingOutput(
+ loss=total_loss,
+ prediction_logits=prediction_scores,
+ seq_relationship_logits=seq_relationship_score,
+ hidden_states=outputs.hidden_states,
+ )
+
+
+@auto_docstring
+class FNetForMaskedLM(FNetPreTrainedModel):
+ _tied_weights_keys = {
+ "cls.predictions.decoder.bias": "cls.predictions.bias",
+ "cls.predictions.decoder.weight": "fnet.embeddings.word_embeddings.weight",
+ }
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.fnet = FNetModel(config)
+ self.cls = FNetOnlyMLMHead(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.cls.predictions.decoder
+
+ def set_output_embeddings(self, new_embeddings):
+ self.cls.predictions.decoder = new_embeddings
+ self.cls.predictions.bias = new_embeddings.bias
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | MaskedLMOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
+ config.vocab_size]` (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]`.
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.fnet(
+ input_ids,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ sequence_output = outputs[0]
+ prediction_scores = self.cls(sequence_output)
+
+ masked_lm_loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss() # -100 index = padding token
+ masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
+
+ if not return_dict:
+ output = (prediction_scores,) + outputs[2:]
+ return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+
+ return MaskedLMOutput(loss=masked_lm_loss, logits=prediction_scores, hidden_states=outputs.hidden_states)
+
+
+@auto_docstring(
+ custom_intro="""
+ FNet Model with a `next sentence prediction (classification)` head on top.
+ """
+)
+class FNetForNextSentencePrediction(FNetPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.fnet = FNetModel(config)
+ self.cls = FNetOnlyNSPHead(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | NextSentencePredictorOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair
+ (see `input_ids` docstring). Indices should be in `[0, 1]`:
+
+ - 0 indicates sequence B is a continuation of sequence A,
+ - 1 indicates sequence B is a random sequence.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, FNetForNextSentencePrediction
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("google/fnet-base")
+ >>> model = FNetForNextSentencePrediction.from_pretrained("google/fnet-base")
+ >>> prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
+ >>> next_sentence = "The sky is blue due to the shorter wavelength of blue light."
+ >>> encoding = tokenizer(prompt, next_sentence, return_tensors="pt")
+ >>> outputs = model(**encoding, labels=torch.LongTensor([1]))
+ >>> logits = outputs.logits
+ >>> assert logits[0, 0] < logits[0, 1] # next sentence was random
+ ```"""
+
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.fnet(
+ input_ids,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ pooled_output = outputs[1]
+
+ seq_relationship_scores = self.cls(pooled_output)
+
+ next_sentence_loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ next_sentence_loss = loss_fct(seq_relationship_scores.view(-1, 2), labels.view(-1))
+
+ if not return_dict:
+ output = (seq_relationship_scores,) + outputs[2:]
+ return ((next_sentence_loss,) + output) if next_sentence_loss is not None else output
+
+ return NextSentencePredictorOutput(
+ loss=next_sentence_loss,
+ logits=seq_relationship_scores,
+ hidden_states=outputs.hidden_states,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ FNet Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooled
+ output) e.g. for GLUE tasks.
+ """
+)
+class FNetForSequenceClassification(FNetPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.fnet = FNetModel(config)
+
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | SequenceClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.fnet(
+ input_ids,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ pooled_output = outputs[1]
+ pooled_output = self.dropout(pooled_output)
+ logits = self.classifier(pooled_output)
+
+ loss = None
+ if labels is not None:
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(logits, labels)
+ if not return_dict:
+ output = (logits,) + outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return SequenceClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states)
+
+
+@auto_docstring
+class FNetForMultipleChoice(FNetPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.fnet = FNetModel(config)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ self.classifier = nn.Linear(config.hidden_size, 1)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | MultipleChoiceModelOutput:
+ r"""
+ input_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary.
+
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+ [`PreTrainedTokenizer.__call__`] for details.
+
+ [What are input IDs?](../glossary#input-ids)
+ token_type_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*):
+ Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,
+ 1]`:
+
+ - 0 corresponds to a *sentence A* token,
+ - 1 corresponds to a *sentence B* token.
+
+ [What are token type IDs?](../glossary#token-type-ids)
+ position_ids (`torch.LongTensor` of shape `(batch_size, num_choices, sequence_length)`, *optional*):
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
+ config.max_position_embeddings - 1]`.
+
+ [What are position IDs?](../glossary#position-ids)
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_choices, 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.
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
+ num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
+ `input_ids` above)
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+ num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
+
+ input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
+ token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
+ position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+ inputs_embeds = (
+ inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
+ if inputs_embeds is not None
+ else None
+ )
+
+ outputs = self.fnet(
+ input_ids,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ pooled_output = outputs[1]
+
+ pooled_output = self.dropout(pooled_output)
+ logits = self.classifier(pooled_output)
+ reshaped_logits = logits.view(-1, num_choices)
+
+ loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(reshaped_logits, labels)
+
+ if not return_dict:
+ output = (reshaped_logits,) + outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return MultipleChoiceModelOutput(loss=loss, logits=reshaped_logits, hidden_states=outputs.hidden_states)
+
+
+@auto_docstring
+class FNetForTokenClassification(FNetPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.fnet = FNetModel(config)
+
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | TokenClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.fnet(
+ input_ids,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ sequence_output = outputs[0]
+
+ sequence_output = self.dropout(sequence_output)
+ logits = self.classifier(sequence_output)
+
+ loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ # Only keep active parts of the loss
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+ if not return_dict:
+ output = (logits,) + outputs[2:]
+ return ((loss,) + output) if loss is not None else output
+
+ return TokenClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states)
+
+
+@auto_docstring
+class FNetForQuestionAnswering(FNetPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.num_labels = config.num_labels
+
+ self.fnet = FNetModel(config)
+ self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ start_positions: torch.Tensor | None = None,
+ end_positions: torch.Tensor | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | QuestionAnsweringModelOutput:
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.fnet(
+ input_ids,
+ token_type_ids=token_type_ids,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ sequence_output = outputs[0]
+
+ logits = self.qa_outputs(sequence_output)
+ start_logits, end_logits = logits.split(1, dim=-1)
+ start_logits = start_logits.squeeze(-1).contiguous()
+ end_logits = end_logits.squeeze(-1).contiguous()
+
+ total_loss = None
+ if start_positions is not None and end_positions is not None:
+ # If we are on multi-GPU, split add a dimension
+ if len(start_positions.size()) > 1:
+ start_positions = start_positions.squeeze(-1)
+ if len(end_positions.size()) > 1:
+ end_positions = end_positions.squeeze(-1)
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
+ ignored_index = start_logits.size(1)
+ start_positions = start_positions.clamp(0, ignored_index)
+ end_positions = end_positions.clamp(0, ignored_index)
+
+ loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+ start_loss = loss_fct(start_logits, start_positions)
+ end_loss = loss_fct(end_logits, end_positions)
+ total_loss = (start_loss + end_loss) / 2
+
+ if not return_dict:
+ output = (start_logits, end_logits) + outputs[2:]
+ return ((total_loss,) + output) if total_loss is not None else output
+
+ return QuestionAnsweringModelOutput(
+ loss=total_loss, start_logits=start_logits, end_logits=end_logits, hidden_states=outputs.hidden_states
+ )
+
+
+__all__ = [
+ "FNetForMaskedLM",
+ "FNetForMultipleChoice",
+ "FNetForNextSentencePrediction",
+ "FNetForPreTraining",
+ "FNetForQuestionAnswering",
+ "FNetForSequenceClassification",
+ "FNetForTokenClassification",
+ "FNetLayer",
+ "FNetModel",
+ "FNetPreTrainedModel",
+]
diff --git a/third_party/transformers/src/transformers/models/fnet/tokenization_fnet.py b/third_party/transformers/src/transformers/models/fnet/tokenization_fnet.py
new file mode 100644
index 0000000000000000000000000000000000000000..20849c147e42a446e7f1537e9b9fe47dd258a1fd
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/fnet/tokenization_fnet.py
@@ -0,0 +1,62 @@
+# Copyright 2024 The HuggingFace Inc. team.
+#
+# 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.
+"""Tokenization classes for FNet model."""
+
+from ...utils import logging
+from ..albert.tokenization_albert import AlbertTokenizer
+
+
+logger = logging.get_logger(__name__)
+
+
+class FNetTokenizer(AlbertTokenizer):
+ """
+ Construct an FNet tokenizer. Based on [Unigram](https://huggingface.co/docs/tokenizers/python/latest/components.html?highlight=unigram#models).
+
+ This tokenizer inherits from [`AlbertTokenizer`] which contains most of the main methods. Users should refer to
+ this superclass for more information regarding those methods.
+
+ Args:
+ do_lower_case (`bool`, *optional*, defaults to `True`):
+ Whether or not to lowercase the input when tokenizing.
+ keep_accents (`bool`, *optional*, defaults to `False`):
+ Whether or not to keep accents when tokenizing.
+ bos_token (`str`, *optional*, defaults to `"[CLS]"`):
+ The beginning of sequence token that was used during pretraining. Can be used a sequence classifier token.
+ eos_token (`str`, *optional*, defaults to `"[SEP]"`):
+ The end of sequence token.
+ unk_token (`str`, *optional*, defaults to `""`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ sep_token (`str`, *optional*, defaults to `"[SEP]"`):
+ The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+ sequence classification or for a text and a question for question answering. It is also used as the last
+ token of a sequence built with special tokens.
+ pad_token (`str`, *optional*, defaults to `""`):
+ The token used for padding, for example when batching sequences of different lengths.
+ cls_token (`str`, *optional*, defaults to `"[CLS]"`):
+ The classifier token which is used when doing sequence classification (classification of the whole sequence
+ instead of per-token classification). It is the first token of the sequence when built with special tokens.
+ mask_token (`str`, *optional*, defaults to `"[MASK]"`):
+ The token used for masking values. This is the token used when training this model with masked language
+ modeling. This is the token which the model will try to predict.
+ """
+
+ model_input_names = ["input_ids", "token_type_ids"]
+
+
+# FNetTokenizerFast is an alias for FNetTokenizer (since AlbertTokenizer is already a fast tokenizer)
+FNetTokenizerFast = FNetTokenizer
+
+__all__ = ["FNetTokenizer", "FNetTokenizerFast"]
diff --git a/third_party/transformers/src/transformers/models/funnel/__init__.py b/third_party/transformers/src/transformers/models/funnel/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e55e71fe491e7cab64dc9c046bab1d60910f34f5
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/funnel/__init__.py
@@ -0,0 +1,29 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_funnel import *
+ from .convert_funnel_original_tf_checkpoint_to_pytorch import *
+ from .modeling_funnel import *
+ from .tokenization_funnel import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/funnel/configuration_funnel.py b/third_party/transformers/src/transformers/models/funnel/configuration_funnel.py
new file mode 100644
index 0000000000000000000000000000000000000000..8ee6a7977545314a05c830e5c0784c7f8764b380
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/funnel/configuration_funnel.py
@@ -0,0 +1,115 @@
+# Copyright 2020, Hugging Face
+#
+# 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.
+"""Funnel Transformer model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="funnel-transformer/small")
+@strict
+class FunnelConfig(PreTrainedConfig):
+ r"""
+ block_sizes (`list[int]`, *optional*, defaults to `[4, 4, 4]`):
+ The sizes of the blocks used in the model.
+ block_repeats (`list[int]`, *optional*):
+ If passed along, each layer of each block is repeated the number of times indicated.
+ num_decoder_layers (`int`, *optional*, defaults to 2):
+ The number of layers in the decoder (when not using the base model).
+ pooling_type (`str`, *optional*, defaults to `"mean"`):
+ Possible values are `"mean"` or `"max"`. The way pooling is performed at the beginning of each block.
+ attention_type (`str`, *optional*, defaults to `"relative_shift"`):
+ Possible values are `"relative_shift"` or `"factorized"`. The former is faster on CPU/GPU while the latter
+ is faster on TPU.
+ separate_cls (`bool`, *optional*, defaults to `True`):
+ Whether or not to separate the cls token when applying pooling.
+ truncate_seq (`bool`, *optional*, defaults to `True`):
+ When using `separate_cls`, whether or not to truncate the last token when pooling, to avoid getting a
+ sequence length that is not a multiple of 2.
+ pool_q_only (`bool`, *optional*, defaults to `True`):
+ Whether or not to apply the pooling only to the query or to query, key and values for the attention layers.
+ """
+
+ model_type = "funnel"
+ attribute_map = {
+ "hidden_size": "d_model",
+ "num_attention_heads": "n_head",
+ }
+
+ vocab_size: int = 30522
+ block_sizes: list[int] | tuple[int, ...] = (4, 4, 4)
+ block_repeats: list[int] | None = None
+ num_decoder_layers: int = 2
+ d_model: int = 768
+ n_head: int = 12
+ d_head: int = 64
+ d_inner: int = 3072
+ hidden_act: str = "gelu_new"
+ hidden_dropout: float | int = 0.1
+ attention_dropout: float | int = 0.1
+ activation_dropout: float | int = 0.0
+ initializer_range: float = 0.1
+ initializer_std: float | None = None
+ layer_norm_eps: float = 1e-9
+ pooling_type: str = "mean"
+ attention_type: str = "relative_shift"
+ separate_cls: bool = True
+ truncate_seq: bool = True
+ pool_q_only: bool = True
+ pad_token_id: int | None = None
+ tie_word_embeddings: bool = True
+
+ def __post_init__(self, **kwargs):
+ self.block_repeats = [1] * len(self.block_sizes) if self.block_repeats is None else self.block_repeats
+ super().__post_init__(**kwargs)
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ if len(self.block_sizes) != len(self.block_repeats):
+ raise ValueError("`block_sizes` and `block_repeats` should have the same length.")
+ if self.pooling_type not in [
+ "mean",
+ "max",
+ ]:
+ raise ValueError(f"Got {self.pooling_type} for `pooling_type` but only 'mean' and 'max' are supported.")
+ if self.attention_type not in [
+ "relative_shift",
+ "factorized",
+ ]:
+ raise ValueError(
+ f"Got {self.attention_type} for `attention_type` but only 'relative_shift' and 'factorized' are supported."
+ )
+
+ @property
+ def num_hidden_layers(self):
+ return sum(self.block_sizes)
+
+ @num_hidden_layers.setter
+ def num_hidden_layers(self, value):
+ raise NotImplementedError(
+ "This model does not support the setting of `num_hidden_layers`. Please set `block_sizes`."
+ )
+
+ @property
+ def num_blocks(self):
+ return len(self.block_sizes)
+
+ @num_blocks.setter
+ def num_blocks(self, value):
+ raise NotImplementedError("This model does not support the setting of `num_blocks`. Please set `block_sizes`.")
+
+
+__all__ = ["FunnelConfig"]
diff --git a/third_party/transformers/src/transformers/models/funnel/convert_funnel_original_tf_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/funnel/convert_funnel_original_tf_checkpoint_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..b16d682f4488d0e35b95a213e221358b3a16f036
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/funnel/convert_funnel_original_tf_checkpoint_to_pytorch.py
@@ -0,0 +1,156 @@
+# Copyright 2020 The HuggingFace Inc. team.
+#
+# 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.
+"""Convert Funnel checkpoint."""
+
+import argparse
+import os
+
+import torch
+
+from transformers import FunnelBaseModel, FunnelConfig, FunnelModel
+from transformers.models.funnel.modeling_funnel import FunnelPositionwiseFFN, FunnelRelMultiheadAttention
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+logging.set_verbosity_info()
+
+
+def load_tf_weights_in_funnel(model, config, tf_checkpoint_path):
+ """Load tf checkpoints in a pytorch model."""
+ try:
+ import re
+
+ import numpy as np
+ import tensorflow as tf
+ except ImportError:
+ logger.error(
+ "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
+ "https://www.tensorflow.org/install/ for installation instructions."
+ )
+ raise
+ tf_path = os.path.abspath(tf_checkpoint_path)
+ logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
+ # Load weights from TF model
+ init_vars = tf.train.list_variables(tf_path)
+ names = []
+ arrays = []
+ for name, shape in init_vars:
+ logger.info(f"Loading TF weight {name} with shape {shape}")
+ array = tf.train.load_variable(tf_path, name)
+ names.append(name)
+ arrays.append(array)
+
+ _layer_map = {
+ "k": "k_head",
+ "q": "q_head",
+ "v": "v_head",
+ "o": "post_proj",
+ "layer_1": "linear_1",
+ "layer_2": "linear_2",
+ "rel_attn": "attention",
+ "ff": "ffn",
+ "kernel": "weight",
+ "gamma": "weight",
+ "beta": "bias",
+ "lookup_table": "weight",
+ "word_embedding": "word_embeddings",
+ "input": "embeddings",
+ }
+
+ for name, array in zip(names, arrays):
+ name = name.split("/")
+ # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
+ # which are not required for using pretrained model
+ if any(
+ n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"]
+ for n in name
+ ):
+ logger.info(f"Skipping {'/'.join(name)}")
+ continue
+ if name[0] == "generator":
+ continue
+ pointer = model
+ skipped = False
+ for m_name in name[1:]:
+ if not isinstance(pointer, FunnelPositionwiseFFN) and re.fullmatch(r"layer_\d+", m_name):
+ layer_index = int(re.search(r"layer_(\d+)", m_name).groups()[0])
+ if layer_index < config.num_hidden_layers:
+ block_idx = 0
+ while layer_index >= config.block_sizes[block_idx]:
+ layer_index -= config.block_sizes[block_idx]
+ block_idx += 1
+ pointer = pointer.blocks[block_idx][layer_index]
+ else:
+ layer_index -= config.num_hidden_layers
+ pointer = pointer.layers[layer_index]
+ elif m_name == "r" and isinstance(pointer, FunnelRelMultiheadAttention):
+ pointer = pointer.r_kernel
+ break
+ elif m_name in _layer_map:
+ pointer = getattr(pointer, _layer_map[m_name])
+ else:
+ try:
+ pointer = getattr(pointer, m_name)
+ except AttributeError:
+ print(f"Skipping {'/'.join(name)}", array.shape)
+ skipped = True
+ break
+ if not skipped:
+ if len(pointer.shape) != len(array.shape):
+ array = array.reshape(pointer.shape)
+ if m_name == "kernel":
+ array = np.transpose(array)
+ pointer.data = torch.from_numpy(array)
+
+ return model
+
+
+def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, config_file, pytorch_dump_path, base_model):
+ # Initialise PyTorch model
+ config = FunnelConfig.from_json_file(config_file)
+ print(f"Building PyTorch model from configuration: {config}")
+ model = FunnelBaseModel(config) if base_model else FunnelModel(config)
+
+ # Load weights from tf checkpoint
+ load_tf_weights_in_funnel(model, config, tf_checkpoint_path)
+
+ # Save pytorch-model
+ print(f"Save PyTorch model to {pytorch_dump_path}")
+ torch.save(model.state_dict(), pytorch_dump_path)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ # Required parameters
+ parser.add_argument(
+ "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
+ )
+ parser.add_argument(
+ "--config_file",
+ default=None,
+ type=str,
+ required=True,
+ help="The config json file corresponding to the pre-trained model. \nThis specifies the model architecture.",
+ )
+ parser.add_argument(
+ "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
+ )
+ parser.add_argument(
+ "--base_model", action="store_true", help="Whether you want just the base model (no decoder) or not."
+ )
+ args = parser.parse_args()
+ convert_tf_checkpoint_to_pytorch(
+ args.tf_checkpoint_path, args.config_file, args.pytorch_dump_path, args.base_model
+ )
diff --git a/third_party/transformers/src/transformers/models/funnel/modeling_funnel.py b/third_party/transformers/src/transformers/models/funnel/modeling_funnel.py
new file mode 100644
index 0000000000000000000000000000000000000000..eb378ebf7b09787269da5d21ec6b88cfec920e66
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/funnel/modeling_funnel.py
@@ -0,0 +1,1364 @@
+# Copyright 2020-present Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
+#
+# 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 Funnel Transformer model."""
+
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...modeling_outputs import (
+ BaseModelOutput,
+ MaskedLMOutput,
+ MultipleChoiceModelOutput,
+ QuestionAnsweringModelOutput,
+ SequenceClassifierOutput,
+ TokenClassifierOutput,
+)
+from ...modeling_utils import PreTrainedModel
+from ...utils import ModelOutput, auto_docstring, logging
+from .configuration_funnel import FunnelConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+INF = 1e6
+
+
+class FunnelEmbeddings(nn.Module):
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__()
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+ self.layer_norm = nn.LayerNorm(config.d_model, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout)
+
+ def forward(
+ self, input_ids: torch.Tensor | None = None, inputs_embeds: torch.Tensor | None = None
+ ) -> torch.Tensor:
+ if inputs_embeds is None:
+ inputs_embeds = self.word_embeddings(input_ids)
+ embeddings = self.layer_norm(inputs_embeds)
+ embeddings = self.dropout(embeddings)
+ return embeddings
+
+
+class FunnelAttentionStructure(nn.Module):
+ """
+ Contains helpers for `FunnelRelMultiheadAttention `.
+ """
+
+ cls_token_type_id: int = 2
+
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__()
+ self.config = config
+ self.sin_dropout = nn.Dropout(config.hidden_dropout)
+ self.cos_dropout = nn.Dropout(config.hidden_dropout)
+ # Track where we are at in terms of pooling from the original input, e.g., by how much the sequence length was
+ # divided.
+ self.pooling_mult = None
+
+ def init_attention_inputs(
+ self,
+ inputs_embeds: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor]:
+ """Returns the attention inputs associated to the inputs of the model."""
+ # inputs_embeds has shape batch_size x seq_len x d_model
+ # attention_mask and token_type_ids have shape batch_size x seq_len
+ self.pooling_mult = 1
+ self.seq_len = seq_len = inputs_embeds.size(1)
+ position_embeds = self.get_position_embeds(seq_len, inputs_embeds.dtype, inputs_embeds.device)
+ token_type_mat = self.token_type_ids_to_mat(token_type_ids) if token_type_ids is not None else None
+ cls_mask = (
+ nn.functional.pad(inputs_embeds.new_ones([seq_len - 1, seq_len - 1]), (1, 0, 1, 0))
+ if self.config.separate_cls
+ else None
+ )
+ return (position_embeds, token_type_mat, attention_mask, cls_mask)
+
+ def token_type_ids_to_mat(self, token_type_ids: torch.Tensor) -> torch.Tensor:
+ """Convert `token_type_ids` to `token_type_mat`."""
+ token_type_mat = token_type_ids[:, :, None] == token_type_ids[:, None]
+ # Treat as in the same segment as both A & B
+ cls_ids = token_type_ids == self.cls_token_type_id
+ cls_mat = cls_ids[:, :, None] | cls_ids[:, None]
+ return cls_mat | token_type_mat
+
+ def get_position_embeds(
+ self, seq_len: int, dtype: torch.dtype, device: torch.device
+ ) -> tuple[torch.Tensor] | list[list[torch.Tensor]]:
+ """
+ Create and cache inputs related to relative position encoding. Those are very different depending on whether we
+ are using the factorized or the relative shift attention:
+
+ For the factorized attention, it returns the matrices (phi, pi, psi, omega) used in the paper, appendix A.2.2,
+ final formula.
+
+ For the relative shift attention, it returns all possible vectors R used in the paper, appendix A.2.1, final
+ formula.
+
+ Paper link: https://huggingface.co/papers/2006.03236
+ """
+ d_model = self.config.d_model
+ if self.config.attention_type == "factorized":
+ # Notations from the paper, appending A.2.2, final formula.
+ # We need to create and return the matrices phi, psi, pi and omega.
+ pos_seq = torch.arange(0, seq_len, 1.0, dtype=torch.int64, device=device).to(dtype)
+ freq_seq = torch.arange(0, d_model // 2, 1.0, dtype=torch.int64, device=device).to(dtype)
+ inv_freq = 1 / (10000 ** (freq_seq / (d_model // 2)))
+ sinusoid = pos_seq[:, None] * inv_freq[None]
+ sin_embed = torch.sin(sinusoid)
+ sin_embed_d = self.sin_dropout(sin_embed)
+ cos_embed = torch.cos(sinusoid)
+ cos_embed_d = self.cos_dropout(cos_embed)
+ # This is different from the formula on the paper...
+ phi = torch.cat([sin_embed_d, sin_embed_d], dim=-1)
+ psi = torch.cat([cos_embed, sin_embed], dim=-1)
+ pi = torch.cat([cos_embed_d, cos_embed_d], dim=-1)
+ omega = torch.cat([-sin_embed, cos_embed], dim=-1)
+ return (phi, pi, psi, omega)
+ else:
+ # Notations from the paper, appending A.2.1, final formula.
+ # We need to create and return all the possible vectors R for all blocks and shifts.
+ freq_seq = torch.arange(0, d_model // 2, 1.0, dtype=torch.int64, device=device).to(dtype)
+ inv_freq = 1 / (10000 ** (freq_seq / (d_model // 2)))
+ # Maximum relative positions for the first input
+ rel_pos_id = torch.arange(-seq_len * 2, seq_len * 2, 1.0, dtype=torch.int64, device=device).to(dtype)
+ zero_offset = seq_len * 2
+ sinusoid = rel_pos_id[:, None] * inv_freq[None]
+ sin_embed = self.sin_dropout(torch.sin(sinusoid))
+ cos_embed = self.cos_dropout(torch.cos(sinusoid))
+ pos_embed = torch.cat([sin_embed, cos_embed], dim=-1)
+
+ pos = torch.arange(0, seq_len, dtype=torch.int64, device=device).to(dtype)
+ pooled_pos = pos
+ position_embeds_list = []
+ for block_index in range(0, self.config.num_blocks):
+ # For each block with block_index > 0, we need two types position embeddings:
+ # - Attention(pooled-q, unpooled-kv)
+ # - Attention(pooled-q, pooled-kv)
+ # For block_index = 0 we only need the second one and leave the first one as None.
+
+ # First type
+ if block_index == 0:
+ position_embeds_pooling = None
+ else:
+ pooled_pos = self.stride_pool_pos(pos, block_index)
+
+ # construct rel_pos_id
+ stride = 2 ** (block_index - 1)
+ rel_pos = self.relative_pos(pos, stride, pooled_pos, shift=2)
+ rel_pos = rel_pos[:, None] + zero_offset
+ rel_pos = rel_pos.expand(rel_pos.size(0), d_model)
+ position_embeds_pooling = torch.gather(pos_embed, 0, rel_pos)
+
+ # Second type
+ pos = pooled_pos
+ stride = 2**block_index
+ rel_pos = self.relative_pos(pos, stride)
+
+ rel_pos = rel_pos[:, None] + zero_offset
+ rel_pos = rel_pos.expand(rel_pos.size(0), d_model)
+ position_embeds_no_pooling = torch.gather(pos_embed, 0, rel_pos)
+
+ position_embeds_list.append([position_embeds_no_pooling, position_embeds_pooling])
+ return position_embeds_list
+
+ def stride_pool_pos(self, pos_id: torch.Tensor, block_index: int):
+ """
+ Pool `pos_id` while keeping the cls token separate (if `config.separate_cls=True`).
+ """
+ if self.config.separate_cls:
+ # Under separate , we treat the as the first token in
+ # the previous block of the 1st real block. Since the 1st real
+ # block always has position 1, the position of the previous block
+ # will be at `1 - 2 ** block_index`.
+ cls_pos = pos_id.new_tensor([-(2**block_index) + 1])
+ pooled_pos_id = pos_id[1:-1] if self.config.truncate_seq else pos_id[1:]
+ return torch.cat([cls_pos, pooled_pos_id[::2]], 0)
+ else:
+ return pos_id[::2]
+
+ def relative_pos(self, pos: torch.Tensor, stride: int, pooled_pos=None, shift: int = 1) -> torch.Tensor:
+ """
+ Build the relative positional vector between `pos` and `pooled_pos`.
+ """
+ if pooled_pos is None:
+ pooled_pos = pos
+
+ ref_point = pooled_pos[0] - pos[0]
+ num_remove = shift * len(pooled_pos)
+ max_dist = ref_point + num_remove * stride
+ min_dist = pooled_pos[0] - pos[-1]
+
+ return torch.arange(max_dist, min_dist - 1, -stride, dtype=torch.long, device=pos.device)
+
+ def stride_pool(
+ self,
+ tensor: torch.Tensor | tuple[torch.Tensor] | list[torch.Tensor],
+ axis: int | tuple[int] | list[int],
+ ) -> torch.Tensor:
+ """
+ Perform pooling by stride slicing the tensor along the given axis.
+ """
+ if tensor is None:
+ return None
+
+ # Do the stride pool recursively if axis is a list or a tuple of ints.
+ if isinstance(axis, (list, tuple)):
+ for ax in axis:
+ tensor = self.stride_pool(tensor, ax)
+ return tensor
+
+ # Do the stride pool recursively if tensor is a list or tuple of tensors.
+ if isinstance(tensor, (tuple, list)):
+ return type(tensor)(self.stride_pool(x, axis) for x in tensor)
+
+ # Deal with negative axis
+ axis %= tensor.ndim
+
+ axis_slice = (
+ slice(None, -1, 2) if self.config.separate_cls and self.config.truncate_seq else slice(None, None, 2)
+ )
+ enc_slice = tuple([slice(None)] * axis + [axis_slice])
+ if self.config.separate_cls:
+ cls_slice = tuple([slice(None)] * axis + [slice(None, 1)])
+ tensor = torch.cat([tensor[cls_slice], tensor], axis=axis)
+ return tensor[enc_slice]
+
+ def pool_tensor(
+ self, tensor: torch.Tensor | tuple[torch.Tensor] | list[torch.Tensor], mode: str = "mean", stride: int = 2
+ ) -> torch.Tensor:
+ """Apply 1D pooling to a tensor of size [B x T (x H)]."""
+ if tensor is None:
+ return None
+
+ # Do the pool recursively if tensor is a list or tuple of tensors.
+ if isinstance(tensor, (tuple, list)):
+ return type(tensor)(self.pool_tensor(tensor, mode=mode, stride=stride) for x in tensor)
+
+ if self.config.separate_cls:
+ suffix = tensor[:, :-1] if self.config.truncate_seq else tensor
+ tensor = torch.cat([tensor[:, :1], suffix], dim=1)
+
+ ndim = tensor.ndim
+ if ndim == 2:
+ tensor = tensor[:, None, :, None]
+ elif ndim == 3:
+ tensor = tensor[:, None, :, :]
+ # Stride is applied on the second-to-last dimension.
+ stride = (stride, 1)
+
+ if mode == "mean":
+ tensor = nn.functional.avg_pool2d(tensor, stride, stride=stride, ceil_mode=True)
+ elif mode == "max":
+ tensor = nn.functional.max_pool2d(tensor, stride, stride=stride, ceil_mode=True)
+ elif mode == "min":
+ tensor = -nn.functional.max_pool2d(-tensor, stride, stride=stride, ceil_mode=True)
+ else:
+ raise NotImplementedError("The supported modes are 'mean', 'max' and 'min'.")
+
+ if ndim == 2:
+ return tensor[:, 0, :, 0]
+ elif ndim == 3:
+ return tensor[:, 0]
+ return tensor
+
+ def pre_attention_pooling(
+ self, output, attention_inputs: tuple[torch.Tensor]
+ ) -> tuple[torch.Tensor, tuple[torch.Tensor]]:
+ """Pool `output` and the proper parts of `attention_inputs` before the attention layer."""
+ position_embeds, token_type_mat, attention_mask, cls_mask = attention_inputs
+ if self.config.pool_q_only:
+ if self.config.attention_type == "factorized":
+ position_embeds = self.stride_pool(position_embeds[:2], 0) + position_embeds[2:]
+ token_type_mat = self.stride_pool(token_type_mat, 1)
+ cls_mask = self.stride_pool(cls_mask, 0)
+ output = self.pool_tensor(output, mode=self.config.pooling_type)
+ else:
+ self.pooling_mult *= 2
+ if self.config.attention_type == "factorized":
+ position_embeds = self.stride_pool(position_embeds, 0)
+ token_type_mat = self.stride_pool(token_type_mat, [1, 2])
+ cls_mask = self.stride_pool(cls_mask, [1, 2])
+ attention_mask = self.pool_tensor(attention_mask, mode="min")
+ output = self.pool_tensor(output, mode=self.config.pooling_type)
+ attention_inputs = (position_embeds, token_type_mat, attention_mask, cls_mask)
+ return output, attention_inputs
+
+ def post_attention_pooling(self, attention_inputs: tuple[torch.Tensor]) -> tuple[torch.Tensor]:
+ """Pool the proper parts of `attention_inputs` after the attention layer."""
+ position_embeds, token_type_mat, attention_mask, cls_mask = attention_inputs
+ if self.config.pool_q_only:
+ self.pooling_mult *= 2
+ if self.config.attention_type == "factorized":
+ position_embeds = position_embeds[:2] + self.stride_pool(position_embeds[2:], 0)
+ token_type_mat = self.stride_pool(token_type_mat, 2)
+ cls_mask = self.stride_pool(cls_mask, 1)
+ attention_mask = self.pool_tensor(attention_mask, mode="min")
+ attention_inputs = (position_embeds, token_type_mat, attention_mask, cls_mask)
+ return attention_inputs
+
+
+def _relative_shift_gather(positional_attn: torch.Tensor, context_len: int, shift: int) -> torch.Tensor:
+ batch_size, n_head, seq_len, max_rel_len = positional_attn.shape
+ # max_rel_len = 2 * context_len + shift -1 is the numbers of possible relative positions i-j
+
+ # What's next is the same as doing the following gather, which might be clearer code but less efficient.
+ # idxs = context_len + torch.arange(0, context_len).unsqueeze(0) - torch.arange(0, seq_len).unsqueeze(1)
+ # # matrix of context_len + i-j
+ # return positional_attn.gather(3, idxs.expand([batch_size, n_head, context_len, context_len]))
+
+ positional_attn = torch.reshape(positional_attn, [batch_size, n_head, max_rel_len, seq_len])
+ positional_attn = positional_attn[:, :, shift:, :]
+ positional_attn = torch.reshape(positional_attn, [batch_size, n_head, seq_len, max_rel_len - shift])
+ positional_attn = positional_attn[..., :context_len]
+ return positional_attn
+
+
+class FunnelRelMultiheadAttention(nn.Module):
+ def __init__(self, config: FunnelConfig, block_index: int) -> None:
+ super().__init__()
+ self.config = config
+ self.block_index = block_index
+ d_model, n_head, d_head = config.d_model, config.n_head, config.d_head
+
+ self.hidden_dropout = nn.Dropout(config.hidden_dropout)
+ self.attention_dropout = nn.Dropout(config.attention_dropout)
+
+ self.q_head = nn.Linear(d_model, n_head * d_head, bias=False)
+ self.k_head = nn.Linear(d_model, n_head * d_head)
+ self.v_head = nn.Linear(d_model, n_head * d_head)
+
+ self.r_w_bias = nn.Parameter(torch.zeros([n_head, d_head]))
+ self.r_r_bias = nn.Parameter(torch.zeros([n_head, d_head]))
+ self.r_kernel = nn.Parameter(torch.zeros([d_model, n_head, d_head]))
+ self.r_s_bias = nn.Parameter(torch.zeros([n_head, d_head]))
+ self.seg_embed = nn.Parameter(torch.zeros([2, n_head, d_head]))
+
+ self.post_proj = nn.Linear(n_head * d_head, d_model)
+ self.layer_norm = nn.LayerNorm(d_model, eps=config.layer_norm_eps)
+ self.scale = 1.0 / (d_head**0.5)
+
+ def relative_positional_attention(self, position_embeds, q_head, context_len, cls_mask=None):
+ """Relative attention score for the positional encodings"""
+ # q_head has shape batch_size x sea_len x n_head x d_head
+ if self.config.attention_type == "factorized":
+ # Notations from the paper, appending A.2.2, final formula (https://huggingface.co/papers/2006.03236)
+ # phi and pi have shape seq_len x d_model, psi and omega have shape context_len x d_model
+ phi, pi, psi, omega = position_embeds
+ # Shape n_head x d_head
+ u = self.r_r_bias * self.scale
+ # Shape d_model x n_head x d_head
+ w_r = self.r_kernel
+
+ # Shape batch_size x sea_len x n_head x d_model
+ q_r_attention = torch.einsum("binh,dnh->bind", q_head + u, w_r)
+ q_r_attention_1 = q_r_attention * phi[:, None]
+ q_r_attention_2 = q_r_attention * pi[:, None]
+
+ # Shape batch_size x n_head x seq_len x context_len
+ positional_attn = torch.einsum("bind,jd->bnij", q_r_attention_1, psi) + torch.einsum(
+ "bind,jd->bnij", q_r_attention_2, omega
+ )
+ else:
+ shift = 2 if q_head.shape[1] != context_len else 1
+ # Notations from the paper, appending A.2.1, final formula (https://huggingface.co/papers/2006.03236)
+ # Grab the proper positional encoding, shape max_rel_len x d_model
+ r = position_embeds[self.block_index][shift - 1]
+ # Shape n_head x d_head
+ v = self.r_r_bias * self.scale
+ # Shape d_model x n_head x d_head
+ w_r = self.r_kernel
+
+ # Shape max_rel_len x n_head x d_model
+ r_head = torch.einsum("td,dnh->tnh", r, w_r)
+ # Shape batch_size x n_head x seq_len x max_rel_len
+ positional_attn = torch.einsum("binh,tnh->bnit", q_head + v, r_head)
+ # Shape batch_size x n_head x seq_len x context_len
+ positional_attn = _relative_shift_gather(positional_attn, context_len, shift)
+
+ if cls_mask is not None:
+ positional_attn *= cls_mask
+ return positional_attn
+
+ def relative_token_type_attention(self, token_type_mat, q_head, cls_mask=None):
+ """Relative attention score for the token_type_ids"""
+ if token_type_mat is None:
+ return 0
+ batch_size, seq_len, context_len = token_type_mat.shape
+ # q_head has shape batch_size x seq_len x n_head x d_head
+ # Shape n_head x d_head
+ r_s_bias = self.r_s_bias * self.scale
+
+ # Shape batch_size x n_head x seq_len x 2
+ token_type_bias = torch.einsum("bind,snd->bnis", q_head + r_s_bias, self.seg_embed)
+ # Shape batch_size x n_head x seq_len x context_len
+ token_type_mat = token_type_mat[:, None].expand([batch_size, q_head.shape[2], seq_len, context_len])
+ # Shapes batch_size x n_head x seq_len
+ diff_token_type, same_token_type = torch.split(token_type_bias, 1, dim=-1)
+ # Shape batch_size x n_head x seq_len x context_len
+ token_type_attn = torch.where(
+ token_type_mat, same_token_type.expand(token_type_mat.shape), diff_token_type.expand(token_type_mat.shape)
+ )
+
+ if cls_mask is not None:
+ token_type_attn *= cls_mask
+ return token_type_attn
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_inputs: tuple[torch.Tensor],
+ output_attentions: bool = False,
+ ) -> tuple[torch.Tensor, ...]:
+ # query has shape batch_size x seq_len x d_model
+ # key and value have shapes batch_size x context_len x d_model
+ position_embeds, token_type_mat, attention_mask, cls_mask = attention_inputs
+
+ batch_size, seq_len, _ = query.shape
+ context_len = key.shape[1]
+ n_head, d_head = self.config.n_head, self.config.d_head
+
+ # Shape batch_size x seq_len x n_head x d_head
+ q_head = self.q_head(query).view(batch_size, seq_len, n_head, d_head)
+ # Shapes batch_size x context_len x n_head x d_head
+ k_head = self.k_head(key).view(batch_size, context_len, n_head, d_head)
+ v_head = self.v_head(value).view(batch_size, context_len, n_head, d_head)
+
+ q_head = q_head * self.scale
+ # Shape n_head x d_head
+ r_w_bias = self.r_w_bias * self.scale
+ # Shapes batch_size x n_head x seq_len x context_len
+ content_score = torch.einsum("bind,bjnd->bnij", q_head + r_w_bias, k_head)
+ positional_attn = self.relative_positional_attention(position_embeds, q_head, context_len, cls_mask)
+ token_type_attn = self.relative_token_type_attention(token_type_mat, q_head, cls_mask)
+
+ # merge attention scores
+ attn_score = content_score + positional_attn + token_type_attn
+
+ # precision safe in case of mixed precision training
+ dtype = attn_score.dtype
+ attn_score = attn_score.float()
+ # perform masking
+ if attention_mask is not None:
+ attn_score = attn_score - INF * (1 - attention_mask[:, None, None].float())
+ # attention probability
+ attn_prob = torch.softmax(attn_score, dim=-1, dtype=dtype)
+ attn_prob = self.attention_dropout(attn_prob)
+
+ # attention output, shape batch_size x seq_len x n_head x d_head
+ attn_vec = torch.einsum("bnij,bjnd->bind", attn_prob, v_head)
+
+ # Shape shape batch_size x seq_len x d_model
+ attn_out = self.post_proj(attn_vec.reshape(batch_size, seq_len, n_head * d_head))
+ attn_out = self.hidden_dropout(attn_out)
+
+ output = self.layer_norm(query + attn_out)
+ return (output, attn_prob) if output_attentions else (output,)
+
+
+class FunnelPositionwiseFFN(nn.Module):
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__()
+ self.linear_1 = nn.Linear(config.d_model, config.d_inner)
+ self.activation_function = ACT2FN[config.hidden_act]
+ self.activation_dropout = nn.Dropout(config.activation_dropout)
+ self.linear_2 = nn.Linear(config.d_inner, config.d_model)
+ self.dropout = nn.Dropout(config.hidden_dropout)
+ self.layer_norm = nn.LayerNorm(config.d_model, config.layer_norm_eps)
+
+ def forward(self, hidden: torch.Tensor) -> torch.Tensor:
+ h = self.linear_1(hidden)
+ h = self.activation_function(h)
+ h = self.activation_dropout(h)
+ h = self.linear_2(h)
+ h = self.dropout(h)
+ return self.layer_norm(hidden + h)
+
+
+class FunnelLayer(nn.Module):
+ def __init__(self, config: FunnelConfig, block_index: int) -> None:
+ super().__init__()
+ self.attention = FunnelRelMultiheadAttention(config, block_index)
+ self.ffn = FunnelPositionwiseFFN(config)
+
+ def forward(
+ self,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_inputs,
+ output_attentions: bool = False,
+ ) -> tuple:
+ attn = self.attention(query, key, value, attention_inputs, output_attentions=output_attentions)
+ output = self.ffn(attn[0])
+ return (output, attn[1]) if output_attentions else (output,)
+
+
+class FunnelEncoder(nn.Module):
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__()
+ self.config = config
+ self.attention_structure = FunnelAttentionStructure(config)
+ self.blocks = nn.ModuleList(
+ [
+ nn.ModuleList([FunnelLayer(config, block_index) for _ in range(block_size)])
+ for block_index, block_size in enumerate(config.block_sizes)
+ ]
+ )
+
+ def forward(
+ self,
+ inputs_embeds: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ output_attentions: bool = False,
+ output_hidden_states: bool = False,
+ return_dict: bool = True,
+ ) -> tuple | BaseModelOutput:
+ # The pooling is not implemented on long tensors, so we convert this mask.
+ attention_mask = attention_mask.type_as(inputs_embeds)
+ attention_inputs = self.attention_structure.init_attention_inputs(
+ inputs_embeds,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ )
+ hidden = inputs_embeds
+
+ all_hidden_states = (inputs_embeds,) if output_hidden_states else None
+ all_attentions = () if output_attentions else None
+
+ for block_index, block in enumerate(self.blocks):
+ pooling_flag = hidden.size(1) > (2 if self.config.separate_cls else 1)
+ pooling_flag = pooling_flag and block_index > 0
+ if pooling_flag:
+ pooled_hidden, attention_inputs = self.attention_structure.pre_attention_pooling(
+ hidden, attention_inputs
+ )
+ for layer_index, layer in enumerate(block):
+ for repeat_index in range(self.config.block_repeats[block_index]):
+ do_pooling = (repeat_index == 0) and (layer_index == 0) and pooling_flag
+ if do_pooling:
+ query = pooled_hidden
+ key = value = hidden if self.config.pool_q_only else pooled_hidden
+ else:
+ query = key = value = hidden
+ layer_output = layer(query, key, value, attention_inputs, output_attentions=output_attentions)
+ hidden = layer_output[0]
+ if do_pooling:
+ attention_inputs = self.attention_structure.post_attention_pooling(attention_inputs)
+
+ if output_attentions:
+ all_attentions = all_attentions + layer_output[1:]
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden,)
+
+ if not return_dict:
+ return tuple(v for v in [hidden, all_hidden_states, all_attentions] if v is not None)
+ return BaseModelOutput(last_hidden_state=hidden, hidden_states=all_hidden_states, attentions=all_attentions)
+
+
+def upsample(
+ x: torch.Tensor, stride: int, target_len: int, separate_cls: bool = True, truncate_seq: bool = False
+) -> torch.Tensor:
+ """
+ Upsample tensor `x` to match `target_len` by repeating the tokens `stride` time on the sequence length dimension.
+ """
+ if stride == 1:
+ return x
+ if separate_cls:
+ cls = x[:, :1]
+ x = x[:, 1:]
+ output = torch.repeat_interleave(x, repeats=stride, dim=1)
+ if separate_cls:
+ if truncate_seq:
+ output = nn.functional.pad(output, (0, 0, 0, stride - 1, 0, 0))
+ output = output[:, : target_len - 1]
+ output = torch.cat([cls, output], dim=1)
+ else:
+ output = output[:, :target_len]
+ return output
+
+
+class FunnelDecoder(nn.Module):
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__()
+ self.config = config
+ self.attention_structure = FunnelAttentionStructure(config)
+ self.layers = nn.ModuleList([FunnelLayer(config, 0) for _ in range(config.num_decoder_layers)])
+
+ def forward(
+ self,
+ final_hidden: torch.Tensor,
+ first_block_hidden: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ output_attentions: bool = False,
+ output_hidden_states: bool = False,
+ return_dict: bool = True,
+ ) -> tuple | BaseModelOutput:
+ upsampled_hidden = upsample(
+ final_hidden,
+ stride=2 ** (len(self.config.block_sizes) - 1),
+ target_len=first_block_hidden.shape[1],
+ separate_cls=self.config.separate_cls,
+ truncate_seq=self.config.truncate_seq,
+ )
+
+ hidden = upsampled_hidden + first_block_hidden
+ all_hidden_states = (hidden,) if output_hidden_states else None
+ all_attentions = () if output_attentions else None
+
+ attention_inputs = self.attention_structure.init_attention_inputs(
+ hidden,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ )
+
+ for layer in self.layers:
+ layer_output = layer(hidden, hidden, hidden, attention_inputs, output_attentions=output_attentions)
+ hidden = layer_output[0]
+
+ if output_attentions:
+ all_attentions = all_attentions + layer_output[1:]
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden,)
+
+ if not return_dict:
+ return tuple(v for v in [hidden, all_hidden_states, all_attentions] if v is not None)
+ return BaseModelOutput(last_hidden_state=hidden, hidden_states=all_hidden_states, attentions=all_attentions)
+
+
+class FunnelDiscriminatorPredictions(nn.Module):
+ """Prediction module for the discriminator, made up of two dense layers."""
+
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__()
+ self.config = config
+ self.dense = nn.Linear(config.d_model, config.d_model)
+ self.dense_prediction = nn.Linear(config.d_model, 1)
+
+ def forward(self, discriminator_hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(discriminator_hidden_states)
+ hidden_states = ACT2FN[self.config.hidden_act](hidden_states)
+ logits = self.dense_prediction(hidden_states).squeeze(-1)
+ return logits
+
+
+@auto_docstring
+class FunnelPreTrainedModel(PreTrainedModel):
+ config: FunnelConfig
+ base_model_prefix = "funnel"
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ classname = module.__class__.__name__
+ if classname.find("Linear") != -1:
+ if getattr(module, "weight", None) is not None:
+ if self.config.initializer_std is None:
+ fan_out, fan_in = module.weight.shape
+ std = np.sqrt(1.0 / float(fan_in + fan_out))
+ else:
+ std = self.config.initializer_std
+ init.normal_(module.weight, std=std)
+ if getattr(module, "bias", None) is not None:
+ init.constant_(module.bias, 0.0)
+ elif classname == "FunnelRelMultiheadAttention":
+ init.uniform_(module.r_w_bias, b=self.config.initializer_range)
+ init.uniform_(module.r_r_bias, b=self.config.initializer_range)
+ init.uniform_(module.r_kernel, b=self.config.initializer_range)
+ init.uniform_(module.r_s_bias, b=self.config.initializer_range)
+ init.uniform_(module.seg_embed, b=self.config.initializer_range)
+ elif classname == "FunnelEmbeddings":
+ std = 1.0 if self.config.initializer_std is None else self.config.initializer_std
+ init.normal_(module.word_embeddings.weight, std=std)
+ if module.word_embeddings.padding_idx is not None:
+ init.zeros_(module.word_embeddings.weight[module.word_embeddings.padding_idx])
+
+
+class FunnelClassificationHead(nn.Module):
+ def __init__(self, config: FunnelConfig, n_labels: int) -> None:
+ super().__init__()
+ self.linear_hidden = nn.Linear(config.d_model, config.d_model)
+ self.dropout = nn.Dropout(config.hidden_dropout)
+ self.linear_out = nn.Linear(config.d_model, n_labels)
+
+ def forward(self, hidden: torch.Tensor) -> torch.Tensor:
+ hidden = self.linear_hidden(hidden)
+ hidden = torch.tanh(hidden)
+ hidden = self.dropout(hidden)
+ return self.linear_out(hidden)
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Output type of [`FunnelForPreTraining`].
+ """
+)
+class FunnelForPreTrainingOutput(ModelOutput):
+ r"""
+ loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
+ Total loss of the ELECTRA-style objective.
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length)`):
+ Prediction scores of the head (scores for each token before SoftMax).
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ The base Funnel Transformer Model transformer outputting raw hidden-states without upsampling head (also called
+ decoder) or any task-specific head on top.
+ """
+)
+class FunnelBaseModel(FunnelPreTrainedModel):
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__(config)
+
+ self.embeddings = FunnelEmbeddings(config)
+ self.encoder = FunnelEncoder(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Embedding:
+ return self.embeddings.word_embeddings
+
+ def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
+ self.embeddings.word_embeddings = new_embeddings
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | BaseModelOutput:
+ 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.return_dict
+
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+ elif input_ids is not None:
+ self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
+ input_shape = input_ids.size()
+ elif inputs_embeds is not None:
+ input_shape = inputs_embeds.size()[:-1]
+ else:
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+ if attention_mask is None:
+ attention_mask = torch.ones(input_shape, device=device)
+ if token_type_ids is None:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+ inputs_embeds = self.embeddings(input_ids, inputs_embeds=inputs_embeds)
+
+ encoder_outputs = self.encoder(
+ inputs_embeds,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ return encoder_outputs
+
+
+@auto_docstring
+class FunnelModel(FunnelPreTrainedModel):
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__(config)
+ self.config = config
+ self.embeddings = FunnelEmbeddings(config)
+ self.encoder = FunnelEncoder(config)
+ self.decoder = FunnelDecoder(config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Embedding:
+ return self.embeddings.word_embeddings
+
+ def set_input_embeddings(self, new_embeddings: nn.Embedding) -> None:
+ self.embeddings.word_embeddings = new_embeddings
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | BaseModelOutput:
+ 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.return_dict
+
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+ elif input_ids is not None:
+ self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
+ input_shape = input_ids.size()
+ elif inputs_embeds is not None:
+ input_shape = inputs_embeds.size()[:-1]
+ else:
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+ if attention_mask is None:
+ attention_mask = torch.ones(input_shape, device=device)
+ if token_type_ids is None:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+ inputs_embeds = self.embeddings(input_ids, inputs_embeds=inputs_embeds)
+
+ encoder_outputs = self.encoder(
+ inputs_embeds,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ output_attentions=output_attentions,
+ output_hidden_states=True,
+ return_dict=return_dict,
+ )
+
+ decoder_outputs = self.decoder(
+ final_hidden=encoder_outputs[0],
+ first_block_hidden=encoder_outputs[1][self.config.block_sizes[0]],
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ if not return_dict:
+ idx = 0
+ outputs = (decoder_outputs[0],)
+ if output_hidden_states:
+ idx += 1
+ outputs = outputs + (encoder_outputs[1] + decoder_outputs[idx],)
+ if output_attentions:
+ idx += 1
+ outputs = outputs + (encoder_outputs[2] + decoder_outputs[idx],)
+ return outputs
+
+ return BaseModelOutput(
+ last_hidden_state=decoder_outputs[0],
+ hidden_states=(encoder_outputs.hidden_states + decoder_outputs.hidden_states)
+ if output_hidden_states
+ else None,
+ attentions=(encoder_outputs.attentions + decoder_outputs.attentions) if output_attentions else None,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Funnel Transformer model with a binary classification head on top as used during pretraining for identifying
+ generated tokens.
+ """
+)
+class FunnelForPreTraining(FunnelPreTrainedModel):
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__(config)
+
+ self.funnel = FunnelModel(config)
+ self.discriminator_predictions = FunnelDiscriminatorPredictions(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | FunnelForPreTrainingOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the ELECTRA-style loss. Input should be a sequence of tokens (see `input_ids`
+ docstring) Indices should be in `[0, 1]`:
+
+ - 0 indicates the token is an original token,
+ - 1 indicates the token was replaced.
+
+ Examples:
+
+ ```python
+ >>> from transformers import AutoTokenizer, FunnelForPreTraining
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("funnel-transformer/small")
+ >>> model = FunnelForPreTraining.from_pretrained("funnel-transformer/small")
+
+ >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+ >>> logits = model(**inputs).logits
+ ```"""
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ discriminator_hidden_states = self.funnel(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ discriminator_sequence_output = discriminator_hidden_states[0]
+
+ logits = self.discriminator_predictions(discriminator_sequence_output)
+
+ loss = None
+ if labels is not None:
+ loss_fct = nn.BCEWithLogitsLoss()
+ if attention_mask is not None:
+ active_loss = attention_mask.view(-1, discriminator_sequence_output.shape[1]) == 1
+ active_logits = logits.view(-1, discriminator_sequence_output.shape[1])[active_loss]
+ active_labels = labels[active_loss]
+ loss = loss_fct(active_logits, active_labels.float())
+ else:
+ loss = loss_fct(logits.view(-1, discriminator_sequence_output.shape[1]), labels.float())
+
+ if not return_dict:
+ output = (logits,) + discriminator_hidden_states[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return FunnelForPreTrainingOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=discriminator_hidden_states.hidden_states,
+ attentions=discriminator_hidden_states.attentions,
+ )
+
+
+@auto_docstring
+class FunnelForMaskedLM(FunnelPreTrainedModel):
+ _tied_weights_keys = {"lm_head.weight": "funnel.embeddings.word_embeddings.weight"}
+
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__(config)
+
+ self.funnel = FunnelModel(config)
+ self.lm_head = nn.Linear(config.d_model, config.vocab_size)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self) -> nn.Linear:
+ return self.lm_head
+
+ def set_output_embeddings(self, new_embeddings: nn.Embedding) -> None:
+ self.lm_head = new_embeddings
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | MaskedLMOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
+ config.vocab_size]` (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]`
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.funnel(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ last_hidden_state = outputs[0]
+ prediction_logits = self.lm_head(last_hidden_state)
+
+ masked_lm_loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss() # -100 index = padding token
+ masked_lm_loss = loss_fct(prediction_logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+ if not return_dict:
+ output = (prediction_logits,) + outputs[1:]
+ return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+
+ return MaskedLMOutput(
+ loss=masked_lm_loss,
+ logits=prediction_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ Funnel Transformer Model with a sequence classification/regression head on top (two linear layer on top of the
+ first timestep of the last hidden state) e.g. for GLUE tasks.
+ """
+)
+class FunnelForSequenceClassification(FunnelPreTrainedModel):
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.config = config
+
+ self.funnel = FunnelBaseModel(config)
+ self.classifier = FunnelClassificationHead(config, config.num_labels)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | SequenceClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.funnel(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ last_hidden_state = outputs[0]
+ pooled_output = last_hidden_state[:, 0]
+ logits = self.classifier(pooled_output)
+
+ loss = None
+ if labels is not None:
+ if self.config.problem_type is None:
+ if self.num_labels == 1:
+ self.config.problem_type = "regression"
+ elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+ self.config.problem_type = "single_label_classification"
+ else:
+ self.config.problem_type = "multi_label_classification"
+
+ if self.config.problem_type == "regression":
+ loss_fct = MSELoss()
+ if self.num_labels == 1:
+ loss = loss_fct(logits.squeeze(), labels.squeeze())
+ else:
+ loss = loss_fct(logits, labels)
+ elif self.config.problem_type == "single_label_classification":
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+ elif self.config.problem_type == "multi_label_classification":
+ loss_fct = BCEWithLogitsLoss()
+ loss = loss_fct(logits, labels)
+
+ if not return_dict:
+ output = (logits,) + outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return SequenceClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring
+class FunnelForMultipleChoice(FunnelPreTrainedModel):
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__(config)
+
+ self.funnel = FunnelBaseModel(config)
+ self.classifier = FunnelClassificationHead(config, 1)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | MultipleChoiceModelOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
+ num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
+ `input_ids` above)
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+ num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]
+
+ input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
+ attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
+ token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
+ inputs_embeds = (
+ inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
+ if inputs_embeds is not None
+ else None
+ )
+
+ outputs = self.funnel(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ last_hidden_state = outputs[0]
+ pooled_output = last_hidden_state[:, 0]
+ logits = self.classifier(pooled_output)
+ reshaped_logits = logits.view(-1, num_choices)
+
+ loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(reshaped_logits, labels)
+
+ if not return_dict:
+ output = (reshaped_logits,) + outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return MultipleChoiceModelOutput(
+ loss=loss,
+ logits=reshaped_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring
+class FunnelForTokenClassification(FunnelPreTrainedModel):
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.funnel = FunnelModel(config)
+ self.dropout = nn.Dropout(config.hidden_dropout)
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ labels: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | TokenClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
+ """
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.funnel(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ last_hidden_state = outputs[0]
+ last_hidden_state = self.dropout(last_hidden_state)
+ logits = self.classifier(last_hidden_state)
+
+ loss = None
+ if labels is not None:
+ loss_fct = CrossEntropyLoss()
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+ if not return_dict:
+ output = (logits,) + outputs[1:]
+ return ((loss,) + output) if loss is not None else output
+
+ return TokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring
+class FunnelForQuestionAnswering(FunnelPreTrainedModel):
+ def __init__(self, config: FunnelConfig) -> None:
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.funnel = FunnelModel(config)
+ self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ token_type_ids: torch.Tensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ start_positions: torch.Tensor | None = None,
+ end_positions: torch.Tensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | QuestionAnsweringModelOutput:
+ return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+ outputs = self.funnel(
+ input_ids,
+ attention_mask=attention_mask,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ last_hidden_state = outputs[0]
+
+ logits = self.qa_outputs(last_hidden_state)
+ start_logits, end_logits = logits.split(1, dim=-1)
+ start_logits = start_logits.squeeze(-1).contiguous()
+ end_logits = end_logits.squeeze(-1).contiguous()
+
+ total_loss = None
+ if start_positions is not None and end_positions is not None:
+ # If we are on multi-GPU, split add a dimension
+ if len(start_positions.size()) > 1:
+ start_positions = start_positions.squeze(-1)
+ if len(end_positions.size()) > 1:
+ end_positions = end_positions.squeeze(-1)
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
+ ignored_index = start_logits.size(1)
+ start_positions = start_positions.clamp(0, ignored_index)
+ end_positions = end_positions.clamp(0, ignored_index)
+
+ loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+ start_loss = loss_fct(start_logits, start_positions)
+ end_loss = loss_fct(end_logits, end_positions)
+ total_loss = (start_loss + end_loss) / 2
+
+ if not return_dict:
+ output = (start_logits, end_logits) + outputs[1:]
+ return ((total_loss,) + output) if total_loss is not None else output
+
+ return QuestionAnsweringModelOutput(
+ loss=total_loss,
+ start_logits=start_logits,
+ end_logits=end_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "FunnelBaseModel",
+ "FunnelForMaskedLM",
+ "FunnelForMultipleChoice",
+ "FunnelForPreTraining",
+ "FunnelForQuestionAnswering",
+ "FunnelForSequenceClassification",
+ "FunnelForTokenClassification",
+ "FunnelModel",
+ "FunnelPreTrainedModel",
+]
diff --git a/third_party/transformers/src/transformers/models/funnel/tokenization_funnel.py b/third_party/transformers/src/transformers/models/funnel/tokenization_funnel.py
new file mode 100644
index 0000000000000000000000000000000000000000..1d1ad42dce52d212762e36ef93477ca650c4a1d8
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/funnel/tokenization_funnel.py
@@ -0,0 +1,164 @@
+# Copyright 2020 The HuggingFace Inc. team.
+#
+# 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.
+"""Tokenization class for Funnel Transformer."""
+
+from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
+from tokenizers.models import WordPiece
+
+from ...tokenization_utils_tokenizers import TokenizersBackend
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
+
+_model_names = [
+ "small",
+ "small-base",
+ "medium",
+ "medium-base",
+ "intermediate",
+ "intermediate-base",
+ "large",
+ "large-base",
+ "xlarge",
+ "xlarge-base",
+]
+
+
+class FunnelTokenizer(TokenizersBackend):
+ r"""
+ Construct a Funnel Transformer tokenizer (backed by HuggingFace's tokenizers library). Based on WordPiece.
+
+ This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should
+ refer to this superclass for more information regarding those methods.
+
+ Args:
+ vocab_file (`str`):
+ File containing the vocabulary.
+ do_lower_case (`bool`, *optional*, defaults to `True`):
+ Whether or not to lowercase the input when tokenizing.
+ unk_token (`str`, *optional*, defaults to `""`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ sep_token (`str`, *optional*, defaults to `""`):
+ The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
+ sequence classification or for a text and a question for question answering. It is also used as the last
+ token of a sequence built with special tokens.
+ pad_token (`str`, *optional*, defaults to `""`):
+ The token used for padding, for example when batching sequences of different lengths.
+ cls_token (`str`, *optional*, defaults to `""`):
+ The classifier token which is used when doing sequence classification (classification of the whole sequence
+ instead of per-token classification). It is the first token of the sequence when built with special tokens.
+ mask_token (`str`, *optional*, defaults to `""`):
+ The token used for masking values. This is the token used when training this model with masked language
+ modeling. This is the token which the model will try to predict.
+ clean_text (`bool`, *optional*, defaults to `True`):
+ Whether or not to clean the text before tokenization by removing any control characters and replacing all
+ whitespaces by the classic one.
+ tokenize_chinese_chars (`bool`, *optional*, defaults to `True`):
+ Whether or not to tokenize Chinese characters. This should likely be deactivated for Japanese (see [this
+ issue](https://github.com/huggingface/transformers/issues/328)).
+ bos_token (`str`, `optional`, defaults to `""`):
+ The beginning of sentence token.
+ eos_token (`str`, `optional`, defaults to `""`):
+ The end of sentence token.
+ strip_accents (`bool`, *optional*):
+ Whether or not to strip all accents. If this option is not specified, then it will be determined by the
+ value for `lowercase` (as in the original BERT).
+ wordpieces_prefix (`str`, *optional*, defaults to `"##"`):
+ The prefix for subwords.
+ vocab (`str` or `dict[str, int]`, *optional*):
+ Custom vocabulary dictionary.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ model = WordPiece
+ cls_token_type_id: int = 2
+
+ def __init__(
+ self,
+ vocab: str | dict[str, int] | None = None,
+ do_lower_case: bool = True,
+ unk_token: str = "",
+ sep_token: str = "",
+ pad_token: str = "",
+ cls_token: str = "",
+ mask_token: str = "",
+ bos_token: str = "",
+ eos_token: str = "",
+ clean_text: bool = True,
+ tokenize_chinese_chars: bool = True,
+ strip_accents: bool | None = None,
+ wordpieces_prefix: str = "##",
+ **kwargs,
+ ):
+ self.do_lower_case = do_lower_case
+ self.tokenize_chinese_chars = tokenize_chinese_chars
+ self.strip_accents = strip_accents
+ self.clean_text = clean_text
+ self.wordpieces_prefix = wordpieces_prefix
+
+ self._vocab = (
+ vocab
+ if vocab is not None
+ else {
+ str(pad_token): 0,
+ str(unk_token): 1,
+ str(cls_token): 2,
+ str(sep_token): 3,
+ str(mask_token): 4,
+ str(bos_token): 5,
+ str(eos_token): 6,
+ }
+ )
+
+ self._tokenizer = Tokenizer(WordPiece(self._vocab, unk_token=str(unk_token)))
+
+ self._tokenizer.normalizer = normalizers.BertNormalizer(
+ clean_text=clean_text,
+ handle_chinese_chars=tokenize_chinese_chars,
+ strip_accents=strip_accents,
+ lowercase=do_lower_case,
+ )
+ self._tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()
+ self._tokenizer.decoder = decoders.WordPiece(prefix=wordpieces_prefix)
+
+ super().__init__(
+ do_lower_case=do_lower_case,
+ unk_token=unk_token,
+ sep_token=sep_token,
+ pad_token=pad_token,
+ cls_token=cls_token,
+ mask_token=mask_token,
+ bos_token=bos_token,
+ eos_token=eos_token,
+ clean_text=clean_text,
+ tokenize_chinese_chars=tokenize_chinese_chars,
+ strip_accents=strip_accents,
+ wordpieces_prefix=wordpieces_prefix,
+ **kwargs,
+ )
+ self._tokenizer.post_processor = processors.TemplateProcessing(
+ single=f"{cls_token}:2 $A:0 {sep_token}:0", # token_type_id is 2 for Funnel transformer
+ pair=f"{cls_token}:2 $A:0 {sep_token}:0 $B:1 {sep_token}:1",
+ special_tokens=[
+ (str(cls_token), self.cls_token_id),
+ (str(sep_token), self.sep_token_id),
+ ],
+ )
+
+
+__all__ = ["FunnelTokenizer"]
diff --git a/third_party/transformers/src/transformers/models/glm4/__init__.py b/third_party/transformers/src/transformers/models/glm4/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6e92a8a2b9c9761d39526ccc0c12c26604fa2a49
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm4/__init__.py
@@ -0,0 +1,27 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_glm4 import *
+ from .modeling_glm4 import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/glm4/configuration_glm4.py b/third_party/transformers/src/transformers/models/glm4/configuration_glm4.py
new file mode 100644
index 0000000000000000000000000000000000000000..f33129607fec2a706251a2aed090ba28ae6851fe
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm4/configuration_glm4.py
@@ -0,0 +1,83 @@
+# Copyright 2025 The GLM4 & ZhipuAI team and HuggingFace Inc. team. All rights reserved.
+#
+#
+# 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.
+
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="zai-org/GLM-OCR")
+@strict
+class Glm4Config(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import Glm4Model, Glm4Config
+ >>> # Initializing a Glm4 glm4-4-9b-chat style configuration
+ >>> configuration = Glm4Config()
+ >>> # Initializing a model from the glm4-4-9b-chat style configuration
+ >>> model = Glm4Model(configuration)
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "glm4"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_proj": "colwise",
+ "layers.*.self_attn.k_proj": "colwise",
+ "layers.*.self_attn.v_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation
+ "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+
+ vocab_size: int = 151552
+ hidden_size: int = 4096
+ intermediate_size: int = 13696
+ num_hidden_layers: int = 40
+ num_attention_heads: int = 32
+ num_key_value_heads: int = 2
+ head_dim: int = 128
+ hidden_act: str = "silu"
+ attention_dropout: float | int = 0.0
+ max_position_embeddings: int = 131072
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 0.00000015625
+ use_cache: bool = True
+ tie_word_embeddings: bool = False
+ rope_parameters: RopeParameters | dict | None = None
+ pad_token_id: int | None = 151329
+ eos_token_id: int | list[int] | None = None
+ bos_token_id: int | None = None
+ attention_bias: bool = True
+
+ def __post_init__(self, **kwargs):
+ kwargs.setdefault("partial_rotary_factor", 0.5) # assign default for BC
+ if self.eos_token_id is None:
+ self.eos_token_id = [151329, 151336, 151338]
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["Glm4Config"]
diff --git a/third_party/transformers/src/transformers/models/glm4/convert_glm4_weights_to_hf.py b/third_party/transformers/src/transformers/models/glm4/convert_glm4_weights_to_hf.py
new file mode 100644
index 0000000000000000000000000000000000000000..01ad00f517ad7c2e27d7426d918c24f633b87ef4
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm4/convert_glm4_weights_to_hf.py
@@ -0,0 +1,199 @@
+import argparse
+import json
+import os
+import re
+
+import torch
+from safetensors.torch import load_file
+from tokenizers import processors
+
+from transformers import Glm4Config, Glm4ForCausalLM, PreTrainedTokenizerFast
+
+
+# fmt: off
+# `None` means we drop the key
+STATE_DICT_MAPPING = {
+ # CausalLM keys
+ r"transformer.output_layer.weight": r"lm_head.weight",
+
+ # Model keys
+ r"transformer.embedding.word_embeddings.weight": r"model.embed_tokens.weight",
+ r"transformer.rotary_pos_emb.inv_freq": None,
+ r"transformer.encoder.final_layernorm.weight": r"model.norm.weight",
+
+ # Layers keys
+ r"transformer.encoder.layers.(\d+).input_layernorm.weight": r"model.layers.\1.input_layernorm.weight",
+
+ # Sandwich keys
+ r"transformer.encoder.layers.(\d+).post_mlp_layernorm.weight": r"model.layers.\1.post_mlp_layernorm.weight",
+ r"transformer.encoder.layers.(\d+).post_self_attn_layernorm.weight": r"model.layers.\1.post_self_attn_layernorm.weight",
+
+ r"transformer.encoder.layers.(\d+).post_attention_layernorm.weight": r"model.layers.\1.post_attention_layernorm.weight",
+
+ # Attention keys
+ r"transformer.encoder.layers.(\d+).self_attention.dense.weight": r"model.layers.\1.self_attn.o_proj.weight",
+ # qkv_proj will later be split in q|k|v|_proj
+ r"transformer.encoder.layers.(\d+).self_attention.query_key_value.(weight|bias)": r"model.layers.\1.self_attn.qkv_proj.\2",
+
+ # MLP keys
+ r"transformer.encoder.layers.(\d+).mlp.dense_h_to_4h.weight": r"model.layers.\1.mlp.gate_up_proj.weight",
+ r"transformer.encoder.layers.(\d+).mlp.dense_4h_to_h.weight": r"model.layers.\1.mlp.down_proj.weight",
+}
+# fmt: on
+
+
+def load_weights(input_dir: str):
+ safetensor_files = [os.path.join(input_dir, x) for x in os.listdir(input_dir) if x.endswith(".safetensors")]
+ bin_files = [os.path.join(input_dir, x) for x in os.listdir(input_dir) if x.endswith(".bin")]
+
+ all_weights = {}
+
+ if safetensor_files:
+ safetensor_files = sorted(safetensor_files, key=lambda x: int(x.rsplit("-", 3)[1]))
+ for file in safetensor_files:
+ tensors = load_file(file)
+ all_weights.update(tensors)
+ return all_weights
+
+ elif bin_files:
+ bin_files = sorted(bin_files, key=lambda x: int(x.rsplit("-", 3)[1]))
+ for file in bin_files:
+ tensors = torch.load(file, map_location="cpu")
+ all_weights.update(tensors)
+ return all_weights
+
+ else:
+ raise ValueError("No .safetensors or .bin files found in the specified directory.")
+
+
+def map_old_key_to_new(old_key):
+ for pattern, replacement in STATE_DICT_MAPPING.items():
+ if replacement is None:
+ if re.fullmatch(pattern, old_key):
+ return None
+ else:
+ new_key, n_replace = re.subn(pattern, replacement, old_key)
+ # Early exit of the loop
+ if n_replace > 0:
+ return new_key
+
+ raise ValueError(f"Key: {old_key} could not be mapped (check the mapping).")
+
+
+def convert_state_dict(original_state_dict: dict, config: Glm4Config):
+ new_dict = {}
+
+ head_dim = config.hidden_size // config.num_attention_heads
+ query_size = config.num_attention_heads * head_dim
+ kv_size = config.num_key_value_heads * head_dim
+
+ for old_key, value in original_state_dict.items():
+ new_key = map_old_key_to_new(old_key)
+ if new_key is None:
+ continue
+
+ if "qkv_proj." in new_key:
+ q_proj, k_proj, v_proj = (
+ value[:query_size, ...],
+ value[query_size : query_size + kv_size, ...],
+ value[query_size + kv_size :, ...],
+ )
+ new_dict[new_key.replace("qkv_proj.", "q_proj.")] = q_proj
+ new_dict[new_key.replace("qkv_proj.", "k_proj.")] = k_proj
+ new_dict[new_key.replace("qkv_proj.", "v_proj.")] = v_proj
+ else:
+ new_dict[new_key] = value
+ return new_dict
+
+
+def convert_config(original_config: dict):
+ key_mapping = {
+ "vocab_size": "padded_vocab_size",
+ "intermediate_size": "ffn_hidden_size",
+ "num_hidden_layers": "num_layers",
+ "max_position_embeddings": "seq_length",
+ "rms_norm_eps": "layernorm_epsilon",
+ "head_dim": "kv_channels",
+ "attention_bias": "add_qkv_bias",
+ }
+ similar_keys_to_keep = [
+ "num_attention_heads",
+ "hidden_size",
+ "attention_dropout",
+ "use_cache",
+ "eos_token_id",
+ "pad_token_id",
+ "tie_word_embeddings",
+ ]
+ new_config_kwargs = {k: original_config[v] for k, v in key_mapping.items()}
+ new_config_kwargs.update({k: v for k, v in original_config.items() if k in similar_keys_to_keep})
+ new_config_kwargs["num_key_value_heads"] = (
+ new_config_kwargs["num_attention_heads"]
+ if not original_config["multi_query_attention"]
+ else original_config["multi_query_group_num"]
+ )
+ new_config_kwargs["rope_theta"] = 10000.0 * getattr(original_config, "rope_ratio", 1)
+
+ new_config = Glm4Config(**new_config_kwargs)
+ return new_config
+
+
+def convert_glm4_tokenizer(input_dir, use_post_processor=False):
+ fast_tok = PreTrainedTokenizerFast.from_pretrained(input_dir, model_input_names=["input_ids", "attention_mask"])
+ if use_post_processor:
+ fast_tok._tokenizer.post_processor = processors.Sequence(
+ [
+ processors.ByteLevel(trim_offsets=False),
+ processors.TemplateProcessing(
+ single="[gMASK]:0 :0 $A:0",
+ pair="[gMASK]:0 :0 $A:0 $B:1",
+ special_tokens=[("[gMASK]", 151331), ("", 151333)],
+ ),
+ ],
+ )
+ else:
+ fast_tok._tokenizer.post_processor = processors.Sequence(
+ [processors.ByteLevel(trim_offsets=False)],
+ )
+ return fast_tok
+
+
+def convert_glm4_model(input_dir, output_dir, use_post_processor=False):
+ # Load and convert config
+ with open(os.path.join(input_dir, "config.json")) as f:
+ original_config = json.load(f)
+ config = convert_config(original_config)
+ config.save_pretrained(output_dir)
+
+ # Load and convert weights
+ original_state_dict = load_weights(input_dir)
+ new_dict = convert_state_dict(original_state_dict, config)
+ with torch.device("meta"):
+ model = Glm4ForCausalLM(config)
+ model.load_state_dict(new_dict, strict=True, assign=True)
+ model.save_pretrained(output_dir)
+
+ # Load and convert tokenizer
+ tokenizer = convert_glm4_tokenizer(input_dir, use_post_processor)
+ tokenizer.save_pretrained(output_dir)
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "input_dir",
+ type=str,
+ help="Location of the local folder copied from the Hub.",
+ )
+ parser.add_argument(
+ "output_dir",
+ type=str,
+ help="Location to write HF model and tokenizer",
+ )
+ parser.add_argument(
+ "--use_post_processor",
+ action="store_true",
+ help="Whether to apply post processor with special tokens",
+ )
+ args = parser.parse_args()
+ convert_glm4_model(args.input_dir, args.output_dir, args.use_post_processor)
diff --git a/third_party/transformers/src/transformers/models/glm4/modeling_glm4.py b/third_party/transformers/src/transformers/models/glm4/modeling_glm4.py
new file mode 100644
index 0000000000000000000000000000000000000000..e99930ae57f620236b2ed8e686f1fbbf53a27d14
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm4/modeling_glm4.py
@@ -0,0 +1,537 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/glm4/modular_glm4.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_glm4.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 The GLM4 & ZhipuAI team and HuggingFace Inc. team. All rights reserved.
+#
+#
+# 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.
+
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import (
+ GenericForSequenceClassification,
+ GenericForTokenClassification,
+ GradientCheckpointingLayer,
+)
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
+from ...utils.generic import maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_glm4 import Glm4Config
+
+
+class Glm4MLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ self.config = config
+ self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
+ self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+ self.activation_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+ up_states = self.gate_up_proj(hidden_states)
+
+ gate, up_states = up_states.chunk(2, dim=-1)
+ up_states = up_states * self.activation_fn(gate)
+
+ return self.down_proj(up_states)
+
+
+class Glm4DecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: Glm4Config, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = Glm4Attention(config=config, layer_idx=layer_idx)
+
+ self.mlp = Glm4MLP(config)
+ self.input_layernorm = Glm4RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = Glm4RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_self_attn_layernorm = Glm4RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_mlp_layernorm = Glm4RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = self.post_self_attn_layernorm(hidden_states)
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = self.post_mlp_layernorm(hidden_states)
+ hidden_states = residual + hidden_states
+ return hidden_states
+
+
+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)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+def rotate_half(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., 0::2]
+ x2 = x[..., 1::2]
+ return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+
+ # Interleave them instead of usual shape
+ cos = cos[..., : cos.shape[-1] // 2].repeat_interleave(2, dim=-1)
+ sin = sin[..., : sin.shape[-1] // 2].repeat_interleave(2, dim=-1)
+
+ # Keep half or full tensor for later concatenation
+ rotary_dim = cos.shape[-1]
+ q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+ k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+ # Apply rotary embeddings on the first half or full tensor
+ q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+ k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+
+ # Concatenate back to full shape
+ q_embed = torch.cat([q_embed, q_pass], dim=-1)
+ k_embed = torch.cat([k_embed, k_pass], dim=-1)
+ return q_embed, k_embed
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class Glm4Attention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: Glm4Config, layer_idx: int | None = None):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = True
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class Glm4RotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: Glm4Config, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: Glm4Config | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
+ head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+ dim = int(head_dim * partial_rotary_factor)
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Glm4RMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ Glm4RMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+@auto_docstring
+class Glm4PreTrainedModel(PreTrainedModel):
+ config: Glm4Config
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["Glm4DecoderLayer"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": Glm4DecoderLayer,
+ "attentions": Glm4Attention,
+ }
+
+
+@auto_docstring
+class Glm4Model(Glm4PreTrainedModel):
+ def __init__(self, config: Glm4Config):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList(
+ [Glm4DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = Glm4RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = Glm4RotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_embeddings=position_embeddings,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class Glm4ForCausalLM(Glm4PreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+ _tp_plan = {"lm_head": "colwise_gather_output"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = Glm4Model(config)
+ self.vocab_size = config.vocab_size
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithPast:
+ r"""
+ 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]`.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, Glm4ForCausalLM
+
+ >>> model = Glm4ForCausalLM.from_pretrained("THUDM/GLM-4-9B-0414")
+ >>> tokenizer = AutoTokenizer.from_pretrained("THUDM/GLM-4-9B-0414")
+
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+ ```"""
+ outputs: BaseModelOutputWithPast = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+class Glm4ForSequenceClassification(GenericForSequenceClassification, Glm4PreTrainedModel):
+ pass
+
+
+class Glm4ForTokenClassification(GenericForTokenClassification, Glm4PreTrainedModel):
+ pass
+
+
+__all__ = [
+ "Glm4PreTrainedModel",
+ "Glm4Model",
+ "Glm4ForCausalLM",
+ "Glm4ForSequenceClassification",
+ "Glm4ForTokenClassification",
+]
diff --git a/third_party/transformers/src/transformers/models/glm4/modular_glm4.py b/third_party/transformers/src/transformers/models/glm4/modular_glm4.py
new file mode 100644
index 0000000000000000000000000000000000000000..4921ac3172d06318894b6ed464e3ad7b47c9fa21
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm4/modular_glm4.py
@@ -0,0 +1,133 @@
+# Copyright 2025 The GLM4 & ZhipuAI team and HuggingFace Inc. team. All rights reserved.
+#
+#
+# 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.
+
+import torch
+
+from ...cache_utils import Cache
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import CausalLMOutputWithPast
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, logging
+from ..glm.modeling_glm import GlmAttention, GlmForCausalLM, GlmForSequenceClassification, GlmForTokenClassification
+from ..phi3.modeling_phi3 import Phi3MLP
+from .configuration_glm4 import Glm4Config
+from .modeling_glm4 import Glm4RMSNorm
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "THUDM/GLM-4-9B-0414"
+
+
+class Glm4MLP(Phi3MLP):
+ pass
+
+
+class Glm4DecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: Glm4Config, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = Glm4Attention(config=config, layer_idx=layer_idx)
+
+ self.mlp = Glm4MLP(config)
+ self.input_layernorm = Glm4RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = Glm4RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_self_attn_layernorm = Glm4RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_mlp_layernorm = Glm4RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = self.post_self_attn_layernorm(hidden_states)
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = self.post_mlp_layernorm(hidden_states)
+ hidden_states = residual + hidden_states
+ return hidden_states
+
+
+class Glm4Attention(GlmAttention):
+ pass
+
+
+class Glm4ForCausalLM(GlmForCausalLM):
+ def forward(
+ self,
+ **super_kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithPast:
+ r"""
+ 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]`.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, Glm4ForCausalLM
+
+ >>> model = Glm4ForCausalLM.from_pretrained("THUDM/GLM-4-9B-0414")
+ >>> tokenizer = AutoTokenizer.from_pretrained("THUDM/GLM-4-9B-0414")
+
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+ ```"""
+ return super().forward(**super_kwargs)
+
+
+class Glm4ForSequenceClassification(GlmForSequenceClassification):
+ pass
+
+
+class Glm4ForTokenClassification(GlmForTokenClassification):
+ pass
+
+
+__all__ = [
+ "Glm4PreTrainedModel", # noqa: F822
+ "Glm4Model", # noqa: F822
+ "Glm4ForCausalLM",
+ "Glm4ForSequenceClassification",
+ "Glm4ForTokenClassification",
+]
diff --git a/third_party/transformers/src/transformers/models/glm4_moe_lite/__init__.py b/third_party/transformers/src/transformers/models/glm4_moe_lite/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..54baf09a1ec8bbfcbcfe7b8fdcb518adb64d5962
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm4_moe_lite/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2025 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_glm4_moe_lite import *
+ from .modeling_glm4_moe_lite import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py b/third_party/transformers/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py
new file mode 100644
index 0000000000000000000000000000000000000000..a9518ed9c5d3a2b1b603c4e6f5d90446f0e84c5a
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm4_moe_lite/configuration_glm4_moe_lite.py
@@ -0,0 +1,119 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_glm4_moe_lite.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="zai-org/GLM-4.5")
+@strict
+class Glm4MoeLiteConfig(PreTrainedConfig):
+ r"""
+ n_group (`int`, *optional*, defaults to 1):
+ Number of groups for routed experts.
+ rope_interleave (`bool`, *optional*, defaults to `True`):
+ Whether to interleave the rotary position embeddings.
+ mlp_layer_types (`list`, *optional*):
+ MLP (Moe vs Dense) pattern for each layer.
+
+ Example:
+
+ ```python
+ >>> from transformers import Glm4MoeLiteModel, Glm4MoeLiteConfig
+
+ >>> # Initializing a Deepseek-V3 style configuration
+ >>> configuration = Glm4MoeLiteConfig()
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "glm4_moe_lite"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_b_proj": "colwise",
+ "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
+ "layers.*.self_attn.kv_b_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
+ "layers.*.mlp.experts.down_proj": "rowwise",
+ "layers.*.mlp.experts": "moe_tp_experts",
+ "layers.*.mlp.gate_proj": "colwise",
+ "layers.*.mlp.up_proj": "colwise",
+ "layers.*.mlp.down_proj": "rowwise",
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+ attribute_map = {
+ "num_local_experts": "n_routed_experts",
+ "head_dim": "qk_rope_head_dim",
+ }
+
+ vocab_size: int = 154880
+ hidden_size: int = 2048
+ intermediate_size: int = 10240
+ moe_intermediate_size: int = 1536
+ num_hidden_layers: int = 47
+ num_attention_heads: int = 20
+ num_key_value_heads: int = 20
+ n_shared_experts: int = 1
+ n_routed_experts: int = 64
+ routed_scaling_factor: float = 1.8
+ kv_lora_rank: int = 512
+ q_lora_rank: int | None = 768
+ qk_rope_head_dim: int = 64
+ v_head_dim: int = 256
+ qk_nope_head_dim: int = 192
+ n_group: int = 1
+ topk_group: int = 1
+ num_experts_per_tok: int = 4
+ norm_topk_prob: bool = True
+ hidden_act: str = "silu"
+ max_position_embeddings: int = 202752
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-5
+ use_cache: bool = True
+ pad_token_id: int | None = None
+ bos_token_id: int | None = 0
+ eos_token_id: int | list[int] | None = 1
+ pretraining_tp: int = 1
+ tie_word_embeddings: bool = False
+ rope_parameters: RopeParameters | dict | None = None
+ rope_interleave: bool = True
+ mlp_layer_types: list[str] | None = None
+ attention_bias: bool = False
+ attention_dropout: float | int = 0.0
+
+ def __post_init__(self, **kwargs):
+ # Default to MoE from the second layer and on
+ if self.mlp_layer_types is None:
+ self.mlp_layer_types = ["dense"] + ["sparse"] * (self.num_hidden_layers - 1)
+ self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["Glm4MoeLiteConfig"]
diff --git a/third_party/transformers/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py b/third_party/transformers/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py
new file mode 100644
index 0000000000000000000000000000000000000000..d59fd2ab996ebe10cd365d8db015a19d3dfb99fb
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm4_moe_lite/modeling_glm4_moe_lite.py
@@ -0,0 +1,725 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_glm4_moe_lite.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+
+import math
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_experts_implementation, use_kernel_forward_from_hub, use_kernel_func_from_hub
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
+from ...utils.generic import is_flash_attention_requested, maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_glm4_moe_lite import Glm4MoeLiteConfig
+
+
+class Glm4MoeLiteRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: Glm4MoeLiteConfig, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: Glm4MoeLiteConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
+ head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+ dim = int(head_dim * partial_rotary_factor)
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+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)
+
+
+@use_kernel_func_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+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)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+def apply_rotary_pos_emb_interleave(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+ r"""
+ TODO let's just use the original freqcis computation to not have the view
+ transpose + reshape! This is not optimized!
+ Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ position_ids (`torch.Tensor`):
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
+ used to pass offsetted position ids when working with a KV-cache.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+
+ b, h, s, d = q.shape
+ q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
+
+ b, h, s, d = k.shape
+ k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
+
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+def yarn_get_mscale(scale=1, mscale=1):
+ if scale <= 1:
+ return 1.0
+ return 0.1 * mscale * math.log(scale) + 1.0
+
+
+class Glm4MoeLiteAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: Glm4MoeLiteConfig, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.attention_dropout = config.attention_dropout
+ self.num_heads = config.num_attention_heads
+
+ self.q_lora_rank = config.q_lora_rank
+ self.qk_rope_head_dim = config.qk_rope_head_dim
+ self.kv_lora_rank = config.kv_lora_rank
+ self.v_head_dim = config.v_head_dim
+ self.qk_nope_head_dim = config.qk_nope_head_dim
+ self.qk_head_dim = config.qk_head_dim
+
+ self.is_causal = True
+ if self.q_lora_rank is None:
+ self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.qk_head_dim, bias=False)
+ else:
+ self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=config.attention_bias)
+ self.q_a_layernorm = Glm4MoeLiteRMSNorm(config.q_lora_rank)
+ self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False)
+
+ self.kv_a_proj_with_mqa = nn.Linear(
+ config.hidden_size,
+ self.kv_lora_rank + self.qk_rope_head_dim,
+ bias=config.attention_bias,
+ )
+ self.kv_a_layernorm = Glm4MoeLiteRMSNorm(self.kv_lora_rank)
+ self.kv_b_proj = nn.Linear(
+ self.kv_lora_rank,
+ self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
+ bias=False,
+ )
+
+ self.o_proj = nn.Linear(
+ self.num_heads * self.v_head_dim,
+ config.hidden_size,
+ bias=config.attention_bias,
+ )
+
+ self.scaling = self.qk_head_dim ** (-0.5)
+ if self.config.rope_parameters.get("rope_type", "default") != "default":
+ mscale_all_dim = self.config.rope_parameters.get("mscale_all_dim", 0)
+ scaling_factor = self.config.rope_parameters["factor"]
+ if mscale_all_dim:
+ mscale = yarn_get_mscale(scaling_factor, mscale_all_dim)
+ self.scaling = self.scaling * mscale * mscale
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: torch.Tensor | None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+ batch_size, seq_length = hidden_states.shape[:-1]
+ query_shape = (batch_size, seq_length, -1, self.qk_head_dim)
+ key_shape = (batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim)
+
+ if self.q_lora_rank is None:
+ q_states = self.q_proj(hidden_states)
+ else:
+ q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
+ q_states = q_states.view(query_shape).transpose(1, 2)
+ q_pass, q_rot = torch.split(q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+
+ compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
+ k_pass, k_rot = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+
+ k_pass = self.kv_b_proj(self.kv_a_layernorm(k_pass)).view(key_shape).transpose(1, 2)
+ k_pass, value_states = torch.split(k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+
+ k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim)
+
+ cos, sin = position_embeddings
+ if self.config.rope_interleave: # support using interleaved weights for efficiency
+ q_rot, k_rot = apply_rotary_pos_emb_interleave(q_rot, k_rot, cos, sin)
+ else:
+ q_rot, k_rot = apply_rotary_pos_emb(q_rot, k_rot, cos, sin)
+ k_rot = k_rot.expand(*k_pass.shape[:-1], -1)
+
+ query_states = torch.cat((q_pass, q_rot), dim=-1)
+ key_states = torch.cat((k_pass, k_rot), dim=-1)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim:
+ value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim])
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim:
+ attn_output = attn_output[:, :, :, : self.v_head_dim]
+
+ attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class Glm4MoeLiteMLP(nn.Module):
+ def __init__(self, config, intermediate_size=None):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size if intermediate_size is None else 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):
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+class Glm4MoeLiteTopkRouter(nn.Module):
+ def __init__(self, config: Glm4MoeLiteConfig):
+ super().__init__()
+ self.config = config
+ self.top_k = config.num_experts_per_tok
+ self.n_routed_experts = config.n_routed_experts
+ self.routed_scaling_factor = config.routed_scaling_factor
+ self.n_group = config.n_group
+ self.topk_group = config.topk_group
+ self.norm_topk_prob = config.norm_topk_prob
+
+ self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size)))
+ self.register_buffer("e_score_correction_bias", torch.zeros((self.n_routed_experts), dtype=torch.float32))
+
+ def forward(self, hidden_states):
+ hidden_states = hidden_states.view(-1, self.config.hidden_size)
+ router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32))
+ return router_logits
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class Glm4MoeLiteRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ Glm4MoeLiteRMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+@use_experts_implementation
+class Glm4MoeLiteNaiveMoe(nn.Module):
+ """Collection of expert weights stored as 3D tensors."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.num_experts = config.num_local_experts
+ self.hidden_dim = config.hidden_size
+ self.intermediate_dim = config.moe_intermediate_size
+ self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
+ self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ top_k_index: torch.Tensor,
+ top_k_weights: torch.Tensor,
+ ) -> torch.Tensor:
+ final_hidden_states = torch.zeros_like(hidden_states)
+ with torch.no_grad():
+ expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
+ expert_mask = expert_mask.permute(2, 1, 0)
+ expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+ for expert_idx in expert_hit:
+ expert_idx = expert_idx[0]
+ if expert_idx == self.num_experts:
+ continue
+ top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+ current_state = hidden_states[token_idx]
+ gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
+ current_hidden_states = self.act_fn(gate) * up
+ current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
+ current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
+ final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
+
+ return final_hidden_states
+
+
+class Glm4MoeLiteMoE(nn.Module):
+ """
+ A mixed expert module containing shared experts.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.experts = Glm4MoeLiteNaiveMoe(config)
+ self.gate = Glm4MoeLiteTopkRouter(config)
+ self.shared_experts = Glm4MoeLiteMLP(
+ config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts
+ )
+ self.n_routed_experts = config.n_routed_experts
+ self.n_group = config.n_group
+ self.topk_group = config.topk_group
+ self.norm_topk_prob = config.norm_topk_prob
+ self.routed_scaling_factor = config.routed_scaling_factor
+ self.top_k = config.num_experts_per_tok
+
+ def route_tokens_to_experts(self, router_logits):
+ router_logits = router_logits.sigmoid()
+ router_logits_for_choice = router_logits + self.gate.e_score_correction_bias
+ group_scores = (
+ router_logits_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group)
+ .topk(2, dim=-1)[0]
+ .sum(dim=-1)
+ )
+ group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+ group_mask = torch.zeros_like(group_scores)
+ group_mask.scatter_(1, group_idx, 1)
+ score_mask = (
+ group_mask.unsqueeze(-1)
+ .expand(-1, self.n_group, self.n_routed_experts // self.n_group)
+ .reshape(-1, self.n_routed_experts)
+ )
+ scores_for_choice = router_logits_for_choice.masked_fill(~score_mask.bool(), 0.0)
+ topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]
+ topk_weights = router_logits.gather(1, topk_indices)
+ if self.norm_topk_prob:
+ denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20
+ topk_weights /= denominator
+ topk_weights = topk_weights * self.routed_scaling_factor
+ return topk_indices, topk_weights
+
+ def forward(self, hidden_states):
+ residuals = hidden_states
+ orig_shape = hidden_states.shape
+ router_logits = self.gate(hidden_states)
+ topk_indices, topk_weights = self.route_tokens_to_experts(router_logits)
+ hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+ hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape)
+ hidden_states = hidden_states + self.shared_experts(residuals)
+ return hidden_states
+
+
+class Glm4MoeLiteDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: Glm4MoeLiteConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = Glm4MoeLiteAttention(config, layer_idx)
+
+ if config.mlp_layer_types[layer_idx] == "sparse":
+ self.mlp = Glm4MoeLiteMoE(config)
+ else:
+ self.mlp = Glm4MoeLiteMLP(config)
+
+ self.input_layernorm = Glm4MoeLiteRMSNorm(config.hidden_size, config.rms_norm_eps)
+ self.post_attention_layernorm = Glm4MoeLiteRMSNorm(config.hidden_size, config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ 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
+ return hidden_states
+
+
+@auto_docstring
+class Glm4MoeLitePreTrainedModel(PreTrainedModel):
+ config: Glm4MoeLiteConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["Glm4MoeLiteDecoderLayer"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": Glm4MoeLiteDecoderLayer,
+ "attentions": Glm4MoeLiteAttention,
+ }
+ _keep_in_fp32_modules_strict = ["e_score_correction_bias"]
+ _keys_to_ignore_on_load_unexpected = [r"model\.layers\.47.*"]
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, Glm4MoeLiteTopkRouter):
+ init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+ init.zeros_(module.e_score_correction_bias)
+ elif isinstance(module, Glm4MoeLiteNaiveMoe):
+ init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range)
+ init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range)
+
+
+@auto_docstring
+class Glm4MoeLiteModel(Glm4MoeLitePreTrainedModel):
+ def __init__(self, config: Glm4MoeLiteConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList(
+ [Glm4MoeLiteDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = Glm4MoeLiteRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = Glm4MoeLiteRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_embeddings=position_embeddings,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class Glm4MoeLiteForCausalLM(Glm4MoeLitePreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+ _tp_plan = {"lm_head": "colwise_gather_output"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = Glm4MoeLiteModel(config)
+ self.vocab_size = config.vocab_size
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> CausalLMOutputWithPast:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, Glm4MoeLiteForCausalLM
+
+ >>> model = Glm4MoeLiteForCausalLM.from_pretrained("meta-glm4_moe_lite/Glm4MoeLite-2-7b-hf")
+ >>> tokenizer = AutoTokenizer.from_pretrained("meta-glm4_moe_lite/Glm4MoeLite-2-7b-hf")
+
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+ ```"""
+ outputs: BaseModelOutputWithPast = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = ["Glm4MoeLitePreTrainedModel", "Glm4MoeLiteModel", "Glm4MoeLiteForCausalLM"]
diff --git a/third_party/transformers/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py b/third_party/transformers/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f65f44a525ae885c7cd612b689f82c74c70a9e4
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm4_moe_lite/modular_glm4_moe_lite.py
@@ -0,0 +1,187 @@
+# Copyright 2025 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+
+import torch.nn as nn
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+from ..deepseek_v3.modeling_deepseek_v3 import DeepseekV3Attention
+from ..glm4_moe.modeling_glm4_moe import (
+ Glm4MoeDecoderLayer,
+ Glm4MoeForCausalLM,
+ Glm4MoeMLP,
+ Glm4MoeModel,
+ Glm4MoeMoE,
+ Glm4MoeNaiveMoe,
+ Glm4MoePreTrainedModel,
+ Glm4MoeRMSNorm,
+ Glm4MoeRotaryEmbedding,
+ Glm4MoeTopkRouter,
+)
+
+
+@auto_docstring(checkpoint="zai-org/GLM-4.5")
+@strict
+class Glm4MoeLiteConfig(PreTrainedConfig):
+ r"""
+ n_group (`int`, *optional*, defaults to 1):
+ Number of groups for routed experts.
+ rope_interleave (`bool`, *optional*, defaults to `True`):
+ Whether to interleave the rotary position embeddings.
+ mlp_layer_types (`list`, *optional*):
+ MLP (Moe vs Dense) pattern for each layer.
+
+ Example:
+
+ ```python
+ >>> from transformers import Glm4MoeLiteModel, Glm4MoeLiteConfig
+
+ >>> # Initializing a Deepseek-V3 style configuration
+ >>> configuration = Glm4MoeLiteConfig()
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "glm4_moe_lite"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_b_proj": "colwise",
+ "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
+ "layers.*.self_attn.kv_b_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
+ "layers.*.mlp.experts.down_proj": "rowwise",
+ "layers.*.mlp.experts": "moe_tp_experts",
+ "layers.*.mlp.gate_proj": "colwise",
+ "layers.*.mlp.up_proj": "colwise",
+ "layers.*.mlp.down_proj": "rowwise",
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+ attribute_map = {
+ "num_local_experts": "n_routed_experts",
+ "head_dim": "qk_rope_head_dim",
+ }
+
+ vocab_size: int = 154880
+ hidden_size: int = 2048
+ intermediate_size: int = 10240
+ moe_intermediate_size: int = 1536
+ num_hidden_layers: int = 47
+ num_attention_heads: int = 20
+ num_key_value_heads: int = 20
+ n_shared_experts: int = 1
+ n_routed_experts: int = 64
+ routed_scaling_factor: float = 1.8
+ kv_lora_rank: int = 512
+ q_lora_rank: int | None = 768
+ qk_rope_head_dim: int = 64
+ v_head_dim: int = 256
+ qk_nope_head_dim: int = 192
+ n_group: int = 1
+ topk_group: int = 1
+ num_experts_per_tok: int = 4
+ norm_topk_prob: bool = True
+ hidden_act: str = "silu"
+ max_position_embeddings: int = 202752
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-5
+ use_cache: bool = True
+ pad_token_id: int | None = None
+ bos_token_id: int | None = 0
+ eos_token_id: int | list[int] | None = 1
+ pretraining_tp: int = 1
+ tie_word_embeddings: bool = False
+ rope_parameters: RopeParameters | dict | None = None
+ rope_interleave: bool = True
+ mlp_layer_types: list[str] | None = None
+ attention_bias: bool = False
+ attention_dropout: float | int = 0.0
+
+ def __post_init__(self, **kwargs):
+ # Default to MoE from the second layer and on
+ if self.mlp_layer_types is None:
+ self.mlp_layer_types = ["dense"] + ["sparse"] * (self.num_hidden_layers - 1)
+ self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+ super().__post_init__(**kwargs)
+
+
+class Glm4MoeLiteRotaryEmbedding(Glm4MoeRotaryEmbedding):
+ pass
+
+
+class Glm4MoeLiteAttention(DeepseekV3Attention):
+ pass
+
+
+class Glm4MoeLiteMLP(Glm4MoeMLP):
+ pass
+
+
+class Glm4MoeLiteTopkRouter(Glm4MoeTopkRouter):
+ pass
+
+
+class Glm4MoeLiteRMSNorm(Glm4MoeRMSNorm):
+ pass
+
+
+class Glm4MoeLiteNaiveMoe(Glm4MoeNaiveMoe):
+ pass
+
+
+class Glm4MoeLiteMoE(Glm4MoeMoE):
+ pass
+
+
+class Glm4MoeLiteDecoderLayer(Glm4MoeDecoderLayer, nn.Module):
+ def __init__(self, config: Glm4MoeLiteConfig, layer_idx: int):
+ nn.Module.__init__(self)
+ self.hidden_size = config.hidden_size
+ self.self_attn = Glm4MoeLiteAttention(config, layer_idx)
+
+ if config.mlp_layer_types[layer_idx] == "sparse":
+ self.mlp = Glm4MoeLiteMoE(config)
+ else:
+ self.mlp = Glm4MoeLiteMLP(config)
+
+ self.input_layernorm = Glm4MoeLiteRMSNorm(config.hidden_size, config.rms_norm_eps)
+ self.post_attention_layernorm = Glm4MoeLiteRMSNorm(config.hidden_size, config.rms_norm_eps)
+
+
+class Glm4MoeLitePreTrainedModel(Glm4MoePreTrainedModel):
+ _keys_to_ignore_on_load_unexpected = [r"model\.layers\.47.*"]
+
+
+class Glm4MoeLiteModel(Glm4MoeModel):
+ pass
+
+
+class Glm4MoeLiteForCausalLM(Glm4MoeForCausalLM):
+ pass
+
+
+__all__ = [
+ "Glm4MoeLiteConfig",
+ "Glm4MoeLitePreTrainedModel",
+ "Glm4MoeLiteModel",
+ "Glm4MoeLiteForCausalLM",
+]
diff --git a/third_party/transformers/src/transformers/models/glm_image/__init__.py b/third_party/transformers/src/transformers/models/glm_image/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..80d6c18e412fe5a10f04966e9fa7c70402b32e2c
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_image/__init__.py
@@ -0,0 +1,31 @@
+# Copyright 2025 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_glm_image import *
+ from .image_processing_glm_image import *
+ from .image_processing_pil_glm_image import *
+ from .modeling_glm_image import *
+ from .processing_glm_image import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/glm_image/configuration_glm_image.py b/third_party/transformers/src/transformers/models/glm_image/configuration_glm_image.py
new file mode 100644
index 0000000000000000000000000000000000000000..fcd302e355606b21ee5df53361e5249c2cbab33b
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_image/configuration_glm_image.py
@@ -0,0 +1,202 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/glm_image/modular_glm_image.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_glm_image.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="zai-org/GLM-Image")
+@strict
+class GlmImageVQVAEConfig(PreTrainedConfig):
+ model_type = "glm_image_vqmodel"
+ base_config_key = "vq_config"
+
+ embed_dim: int = 2048
+ num_embeddings: int = 16384
+ latent_channels: int = 1536
+ in_channels: int = 3
+ initializer_range: float = 0.02
+
+
+@auto_docstring(checkpoint="zai-org/GLM-Image")
+@strict
+class GlmImageVisionConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import GlmImageVisionConfig, GlmImageVisionModel
+
+ >>> # Initializing a GlmImageVisionConfig GLM-Image style configuration
+ >>> configuration = GlmImageVisionConfig()
+
+ >>> # Initializing a model (with random weights) from the GLM-Image configuration
+ >>> model = GlmImageVisionModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "glm_image_vision"
+ base_config_key = "vision_config"
+
+ depth: int = 40
+ hidden_size: int = 1536
+ hidden_act: str = "gelu"
+ attention_bias: bool = True
+ attention_dropout: float | int = 0.0
+ num_heads: int = 16
+ in_channels: int = 3
+ image_size: int | list[int] | tuple[int, int] = 2048
+ patch_size: int | list[int] | tuple[int, int] = 16
+ spatial_merge_size: int = 1
+ intermediate_size: int = 6144
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-06
+
+
+@auto_docstring(checkpoint="zai-org/GLM-Image")
+@strict
+class GlmImageTextConfig(PreTrainedConfig):
+ r"""
+ vision_vocab_size (`int`, *optional*, defaults to 16512):
+ Vision vocabulary size of the GlmImage model. Defines the number of different tokens that can be
+ represented by the `inputs_ids` passed when calling [`GlmImageVisionModel`]
+
+ Example:
+
+ ```python
+ >>> from transformers import GlmImageTextModel, GlmImageConfig
+
+ >>> # Initializing a GlmImageConfig style configuration
+ >>> configuration = GlmImageConfig()
+
+ >>> # Initializing a model from the GlmImageConfig style configuration
+ >>> model = GlmImageTextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "glm_image_text"
+ base_config_key = "text_config"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ # Default tensor parallel plan for base model `GlmImage`
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_proj": "colwise",
+ "layers.*.self_attn.k_proj": "colwise",
+ "layers.*.self_attn.v_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation
+ "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+ ignore_keys_at_rope_validation = {"mrope_section"}
+
+ vocab_size: int = 168064
+ hidden_size: int = 4096
+ intermediate_size: int = 13696
+ num_hidden_layers: int = 40
+ num_attention_heads: int = 32
+ num_key_value_heads: int | None = 2
+ hidden_act: str = "silu"
+ max_position_embeddings: int = 131072
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-05
+ use_cache: bool = True
+ attention_dropout: float | int = 0.0
+ rope_parameters: RopeParameters | dict | None = None
+ pad_token_id: int = 167841
+ vision_vocab_size: int = 16512
+ attention_bias: bool = True
+ eos_token_id: int | list[int] | None = 16385
+
+ def __post_init__(self, **kwargs):
+ if self.num_key_value_heads is None:
+ self.num_key_value_heads = self.num_attention_heads
+
+ super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="zai-org/GLM-Image")
+@strict
+class GlmImageConfig(PreTrainedConfig):
+ r"""
+ image_start_token_id (`int`, *optional*, defaults to 16384):
+ The image start token index to encode the start of image.
+ image_end_token_id (`int`, *optional*, defaults to 16385):
+ The image end token index to encode the end of image.
+
+ ```python
+ >>> from transformers import Glm4vForConditionalGeneration, Glm4vConfig
+
+ >>> # Initializing a GLM-Image style configuration
+ >>> configuration = Glm4vConfig()
+
+ >>> # Initializing a model from the GLM-Image style configuration
+ >>> model = Glm4vForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "glm_image"
+ sub_configs = {
+ "vision_config": GlmImageVisionConfig,
+ "text_config": GlmImageTextConfig,
+ "vq_config": GlmImageVQVAEConfig,
+ }
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ text_config: dict | PreTrainedConfig | None = None
+ vision_config: dict | PreTrainedConfig | None = None
+ vq_config: dict | PreTrainedConfig | None = None
+ image_token_id: int = 167855
+ image_start_token_id: int = 16384
+ image_end_token_id: int = 16385
+ tie_word_embeddings: bool = False
+
+ def __post_init__(self, **kwargs):
+ if isinstance(self.vision_config, dict):
+ self.vision_config = self.sub_configs["vision_config"](**self.vision_config)
+ elif self.vision_config is None:
+ self.vision_config = self.sub_configs["vision_config"](**kwargs)
+
+ if isinstance(self.vq_config, dict):
+ self.vq_config = self.sub_configs["vq_config"](**self.vq_config)
+ elif self.vq_config is None:
+ self.vq_config = self.sub_configs["vq_config"](**kwargs)
+
+ if isinstance(self.text_config, dict):
+ self.text_config = self.sub_configs["text_config"](**self.text_config)
+ elif self.text_config is None:
+ self.text_config = self.sub_configs["text_config"](**kwargs)
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["GlmImageVQVAEConfig", "GlmImageVisionConfig", "GlmImageTextConfig", "GlmImageConfig"]
diff --git a/third_party/transformers/src/transformers/models/glm_image/image_processing_glm_image.py b/third_party/transformers/src/transformers/models/glm_image/image_processing_glm_image.py
new file mode 100644
index 0000000000000000000000000000000000000000..6412ad5d3a481dfb4496ee565ffc0b5e95d152cf
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_image/image_processing_glm_image.py
@@ -0,0 +1,270 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/glm_image/modular_glm_image.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_glm_image.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+import math
+from collections.abc import Iterable
+
+import torch
+from torchvision.transforms.v2 import functional as tvF
+
+from ...feature_extraction_utils import BatchFeature
+from ...image_processing_backends import TorchvisionBackend
+from ...image_transforms import group_images_by_shape, reorder_images
+from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, ImageInput, PILImageResampling, SizeDict
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring
+
+
+class GlmImageImageProcessorKwargs(ImagesKwargs, total=False):
+ r"""
+ min_pixels (`int`, *optional*, defaults to `56 * 56`):
+ The min pixels of the image to resize the image.
+ max_pixels (`int`, *optional*, defaults to `28 * 28 * 1280`):
+ The max pixels of the image to resize the image.
+ patch_size (`int`, *optional*, defaults to 14):
+ The spatial patch size of the vision encoder.
+ temporal_patch_size (`int`, *optional*, defaults to 2):
+ The temporal patch size of the vision encoder.
+ merge_size (`int`, *optional*, defaults to 2):
+ The merge size of the vision encoder to llm encoder.
+ """
+
+ min_pixels: int
+ max_pixels: int
+ patch_size: int
+ temporal_patch_size: int
+ merge_size: int
+
+
+def smart_resize(
+ height: int,
+ width: int,
+ factor: int = 16,
+ min_pixels: int = 512 * 512,
+ max_pixels: int = 2048 * 2048,
+) -> tuple[int, int]:
+ if height < factor or width < factor:
+ raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}")
+ elif max(height, width) / min(height, width) > 4:
+ raise ValueError(
+ f"absolute aspect ratio must be smaller than 4, got {max(height, width) / min(height, width)}"
+ )
+
+ shortest_edge = int(round(math.sqrt(min_pixels)))
+ longest_edge = int(round(math.sqrt(max_pixels)))
+ min_side = min(height, width)
+ max_side = max(height, width)
+
+ scale = 1.0
+
+ if min_side < shortest_edge:
+ scale = shortest_edge / min_side
+
+ if max_side * scale > longest_edge:
+ scale = longest_edge / max_side
+
+ height = height // 2
+ width = width // 2
+
+ h_bar = max(factor, int(round(height * scale / factor)) * factor)
+ w_bar = max(factor, int(round(width * scale / factor)) * factor)
+
+ if max(h_bar, w_bar) > longest_edge:
+ beta = max(h_bar, w_bar) / longest_edge
+ h_bar = max(factor, int(math.floor((h_bar / beta) / factor)) * factor)
+ w_bar = max(factor, int(math.floor((w_bar / beta) / factor)) * factor)
+
+ return h_bar, w_bar
+
+
+@auto_docstring
+class GlmImageImageProcessor(TorchvisionBackend):
+ do_resize = True
+ resample = PILImageResampling.BICUBIC
+ size = {"shortest_edge": 56 * 56, "longest_edge": 28 * 28 * 1280}
+ default_to_square = False
+ do_rescale = True
+ do_normalize = True
+ image_mean = OPENAI_CLIP_MEAN
+ image_std = OPENAI_CLIP_STD
+ do_convert_rgb = True
+ patch_size = 14
+ temporal_patch_size = 2
+ merge_size = 2
+ valid_kwargs = GlmImageImageProcessorKwargs
+ model_input_names = ["pixel_values", "image_grid_thw", "images_per_sample"]
+
+ def __init__(self, **kwargs: Unpack[GlmImageImageProcessorKwargs]):
+ size = kwargs.pop("size", None)
+ min_pixels = kwargs.pop("min_pixels", None)
+ max_pixels = kwargs.pop("max_pixels", None)
+ # backward compatibility: override size with min_pixels and max_pixels if they are provided
+ size = self.size if size is None else size
+ if min_pixels is not None:
+ size["shortest_edge"] = min_pixels
+ size.pop("min_pixels", None)
+ if max_pixels is not None:
+ size["longest_edge"] = max_pixels
+ size.pop("max_pixels", None)
+ if "shortest_edge" not in size or "longest_edge" not in size:
+ raise ValueError("size must contain 'shortest_edge' and 'longest_edge' keys.")
+
+ super().__init__(size=size, **kwargs)
+
+ def _standardize_kwargs(
+ self,
+ size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
+ min_pixels: int | None = None,
+ max_pixels: int | None = None,
+ **kwargs,
+ ) -> dict:
+ if min_pixels is not None and max_pixels is not None:
+ size = SizeDict(shortest_edge=min_pixels, longest_edge=max_pixels)
+ kwargs = super()._standardize_kwargs(size=size, **kwargs)
+ size = kwargs.get("size", self.size)
+ if not size.shortest_edge or not size.longest_edge:
+ raise ValueError("size must contain 'shortest_edge' and 'longest_edge' keys.")
+ return kwargs
+
+ @auto_docstring
+ def preprocess(
+ self,
+ images: ImageInput,
+ **kwargs: Unpack[GlmImageImageProcessorKwargs],
+ ) -> BatchFeature:
+ return super().preprocess(images, **kwargs)
+
+ def _preprocess(
+ self,
+ images: list["torch.Tensor"],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | tvF.InterpolationMode | int | None",
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ patch_size: int,
+ temporal_patch_size: int,
+ merge_size: int,
+ disable_grouping: bool | None,
+ return_tensors: str | TensorType | None,
+ **kwargs,
+ ) -> BatchFeature:
+ grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping)
+ resized_images_grouped = {}
+ for shape, stacked_images in grouped_images.items():
+ height, width = stacked_images.shape[-2:]
+ if do_resize:
+ resized_height, resized_width = smart_resize(
+ height,
+ width,
+ factor=patch_size * merge_size,
+ min_pixels=size.shortest_edge,
+ max_pixels=size.longest_edge,
+ )
+ stacked_images = self.resize(
+ image=stacked_images,
+ size=SizeDict(height=resized_height, width=resized_width),
+ resample=resample,
+ )
+ resized_images_grouped[shape] = stacked_images
+ resized_images = reorder_images(resized_images_grouped, grouped_images_index)
+
+ grouped_images, grouped_images_index = group_images_by_shape(resized_images, disable_grouping=disable_grouping)
+ processed_images_grouped = {}
+ processed_grids = {}
+ for shape, stacked_images in grouped_images.items():
+ resized_height, resized_width = stacked_images.shape[-2:]
+ patches = self.rescale_and_normalize(
+ stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std
+ )
+ if patches.ndim == 4:
+ patches = patches.unsqueeze(1)
+ if patches.shape[1] % temporal_patch_size != 0:
+ repeats = patches[:, -1:].repeat(1, temporal_patch_size - 1, 1, 1, 1)
+ patches = torch.cat([patches, repeats], dim=1)
+ batch_size, grid_t, channel = patches.shape[:3]
+ grid_t = grid_t // temporal_patch_size
+ grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+ patches = patches.view(
+ batch_size,
+ grid_t,
+ temporal_patch_size,
+ channel,
+ grid_h // merge_size,
+ merge_size,
+ patch_size,
+ grid_w // merge_size,
+ merge_size,
+ patch_size,
+ )
+ patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)
+ flatten_patches = patches.reshape(
+ batch_size,
+ grid_t * grid_h * grid_w,
+ channel * temporal_patch_size * patch_size * patch_size,
+ )
+
+ processed_images_grouped[shape] = flatten_patches
+ processed_grids[shape] = [[grid_t, grid_h, grid_w]] * batch_size
+
+ processed_images = reorder_images(processed_images_grouped, grouped_images_index)
+ processed_grids_ordered = reorder_images(processed_grids, grouped_images_index)
+ pixel_values = torch.cat(processed_images, dim=0)
+ image_grid_thw = torch.tensor(processed_grids_ordered, dtype=torch.long)
+
+ return BatchFeature(
+ data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, tensor_type=return_tensors
+ )
+
+ def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None):
+ """
+ A utility that returns number of image patches for a given image size.
+
+ Note: Do not remove this method! It is used by vLLM to infer the number of patches and placeholders
+ without an image input.
+
+ Args:
+ height (`int`):
+ Height of the input image.
+ width (`int`):
+ Width of the input image.
+ images_kwargs (`dict`, *optional*)
+ Any kwargs to override defaults of the image processor.
+ Returns:
+ `int`: Number of image patches per image.
+ """
+ min_pixels = images_kwargs["min_pixels"] if "min_pixels" in images_kwargs else self.size["shortest_edge"]
+ max_pixels = images_kwargs["max_pixels"] if "max_pixels" in images_kwargs else self.size["longest_edge"]
+ patch_size = images_kwargs.get("patch_size", self.patch_size)
+ merge_size = images_kwargs.get("merge_size", self.merge_size)
+
+ factor = patch_size * merge_size
+ resized_height, resized_width = smart_resize(
+ height, width, factor, min_pixels=min_pixels, max_pixels=max_pixels
+ )
+ grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+ return grid_h * grid_w
+
+
+__all__ = ["GlmImageImageProcessor"]
diff --git a/third_party/transformers/src/transformers/models/glm_image/image_processing_pil_glm_image.py b/third_party/transformers/src/transformers/models/glm_image/image_processing_pil_glm_image.py
new file mode 100644
index 0000000000000000000000000000000000000000..2dde18ef2066332bfd31cab9ca0698e840d1cee4
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_image/image_processing_pil_glm_image.py
@@ -0,0 +1,266 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/glm_image/modular_glm_image.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_glm_image.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+import math
+from collections.abc import Iterable
+
+import numpy as np
+
+from ...feature_extraction_utils import BatchFeature
+from ...image_processing_backends import PilBackend
+from ...image_utils import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD, ImageInput, PILImageResampling, SizeDict
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring
+
+
+# Adapted from transformers.models.glm_image.image_processing_glm_image.GlmImageImageProcessorKwargs
+class GlmImageImageProcessorKwargs(ImagesKwargs, total=False):
+ r"""
+ min_pixels (`int`, *optional*, defaults to `56 * 56`):
+ The min pixels of the image to resize the image.
+ max_pixels (`int`, *optional*, defaults to `28 * 28 * 1280`):
+ The max pixels of the image to resize the image.
+ patch_size (`int`, *optional*, defaults to 14):
+ The spatial patch size of the vision encoder.
+ temporal_patch_size (`int`, *optional*, defaults to 2):
+ The temporal patch size of the vision encoder.
+ merge_size (`int`, *optional*, defaults to 2):
+ The merge size of the vision encoder to llm encoder.
+ """
+
+ min_pixels: int
+ max_pixels: int
+ patch_size: int
+ temporal_patch_size: int
+ merge_size: int
+
+
+def smart_resize(
+ height: int,
+ width: int,
+ factor: int = 16,
+ min_pixels: int = 512 * 512,
+ max_pixels: int = 2048 * 2048,
+) -> tuple[int, int]:
+ if height < factor or width < factor:
+ raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}")
+ elif max(height, width) / min(height, width) > 4:
+ raise ValueError(
+ f"absolute aspect ratio must be smaller than 4, got {max(height, width) / min(height, width)}"
+ )
+
+ shortest_edge = int(round(math.sqrt(min_pixels)))
+ longest_edge = int(round(math.sqrt(max_pixels)))
+ min_side = min(height, width)
+ max_side = max(height, width)
+
+ scale = 1.0
+
+ if min_side < shortest_edge:
+ scale = shortest_edge / min_side
+
+ if max_side * scale > longest_edge:
+ scale = longest_edge / max_side
+
+ height = height // 2
+ width = width // 2
+
+ h_bar = max(factor, int(round(height * scale / factor)) * factor)
+ w_bar = max(factor, int(round(width * scale / factor)) * factor)
+
+ if max(h_bar, w_bar) > longest_edge:
+ beta = max(h_bar, w_bar) / longest_edge
+ h_bar = max(factor, int(math.floor((h_bar / beta) / factor)) * factor)
+ w_bar = max(factor, int(math.floor((w_bar / beta) / factor)) * factor)
+
+ return h_bar, w_bar
+
+
+@auto_docstring
+class GlmImageImageProcessorPil(PilBackend):
+ do_resize = True
+ resample = PILImageResampling.BICUBIC
+ size = {"shortest_edge": 56 * 56, "longest_edge": 28 * 28 * 1280}
+ default_to_square = False
+ do_rescale = True
+ do_normalize = True
+ image_mean = OPENAI_CLIP_MEAN
+ image_std = OPENAI_CLIP_STD
+ do_convert_rgb = True
+ patch_size = 14
+ temporal_patch_size = 2
+ merge_size = 2
+ valid_kwargs = GlmImageImageProcessorKwargs
+ model_input_names = ["pixel_values", "image_grid_thw", "images_per_sample"]
+
+ def __init__(self, **kwargs: Unpack[GlmImageImageProcessorKwargs]):
+ size = kwargs.pop("size", None)
+ min_pixels = kwargs.pop("min_pixels", None)
+ max_pixels = kwargs.pop("max_pixels", None)
+ # backward compatibility: override size with min_pixels and max_pixels if they are provided
+ size = self.size if size is None else size
+ if min_pixels is not None:
+ size["shortest_edge"] = min_pixels
+ size.pop("min_pixels", None)
+ if max_pixels is not None:
+ size["longest_edge"] = max_pixels
+ size.pop("max_pixels", None)
+ if "shortest_edge" not in size or "longest_edge" not in size:
+ raise ValueError("size must contain 'shortest_edge' and 'longest_edge' keys.")
+
+ super().__init__(size=size, **kwargs)
+
+ def _standardize_kwargs(
+ self,
+ size: int | Iterable[int] | dict[str, int] | SizeDict | None = None,
+ min_pixels: int | None = None,
+ max_pixels: int | None = None,
+ **kwargs,
+ ) -> dict:
+ if min_pixels is not None and max_pixels is not None:
+ size = SizeDict(shortest_edge=min_pixels, longest_edge=max_pixels)
+ kwargs = super()._standardize_kwargs(size=size, **kwargs)
+ size = kwargs.get("size", self.size)
+ if not size.shortest_edge or not size.longest_edge:
+ raise ValueError("size must contain 'shortest_edge' and 'longest_edge' keys.")
+ return kwargs
+
+ @auto_docstring
+ def preprocess(
+ self,
+ images: ImageInput,
+ **kwargs: Unpack[GlmImageImageProcessorKwargs],
+ ) -> BatchFeature:
+ return super().preprocess(images, **kwargs)
+
+ def _preprocess(
+ self,
+ images: list[np.ndarray],
+ do_resize: bool,
+ size: SizeDict,
+ resample: "PILImageResampling | None",
+ do_rescale: bool,
+ rescale_factor: float,
+ do_normalize: bool,
+ image_mean: float | list[float] | None,
+ image_std: float | list[float] | None,
+ patch_size: int,
+ temporal_patch_size: int,
+ merge_size: int,
+ return_tensors: str | TensorType | None,
+ **kwargs,
+ ) -> BatchFeature:
+ all_patches = []
+ all_grids = []
+
+ for image in images:
+ height, width = image.shape[-2:]
+ if do_resize:
+ resized_height, resized_width = smart_resize(
+ height,
+ width,
+ factor=patch_size * merge_size,
+ min_pixels=size.shortest_edge,
+ max_pixels=size.longest_edge,
+ )
+ image = self.resize(
+ image,
+ size=SizeDict(height=resized_height, width=resized_width),
+ resample=resample,
+ )
+ else:
+ resized_height, resized_width = height, width
+
+ if do_rescale:
+ image = self.rescale(image, rescale_factor)
+ if do_normalize:
+ image = self.normalize(image, image_mean, image_std)
+
+ patches = np.expand_dims(image, axis=0)
+ if patches.ndim == 4:
+ patches = np.expand_dims(patches, axis=1)
+ if patches.shape[1] % temporal_patch_size != 0:
+ repeats = np.repeat(
+ patches[:, -1:], temporal_patch_size - (patches.shape[1] % temporal_patch_size), axis=1
+ )
+ patches = np.concatenate([patches, repeats], axis=1)
+
+ batch_size, grid_t, channel = patches.shape[0], patches.shape[1] // temporal_patch_size, patches.shape[2]
+ grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+
+ patches = patches.reshape(
+ batch_size,
+ grid_t,
+ temporal_patch_size,
+ channel,
+ grid_h // merge_size,
+ merge_size,
+ patch_size,
+ grid_w // merge_size,
+ merge_size,
+ patch_size,
+ )
+ patches = patches.transpose(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)
+ flatten_patches = patches.reshape(
+ batch_size * grid_t * grid_h * grid_w,
+ channel * temporal_patch_size * patch_size * patch_size,
+ )
+
+ all_patches.append(flatten_patches)
+ all_grids.append([grid_t, grid_h, grid_w])
+
+ pixel_values = np.concatenate(all_patches, axis=0)
+ image_grid_thw = np.array(all_grids, dtype=np.int64)
+
+ return BatchFeature(
+ data={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, tensor_type=return_tensors
+ )
+
+ def get_number_of_image_patches(self, height: int, width: int, images_kwargs=None):
+ """
+ A utility that returns number of image patches for a given image size.
+
+ Note: Do not remove this method! It is used by vLLM to infer the number of patches and placeholders
+ without an image input.
+
+ Args:
+ height (`int`):
+ Height of the input image.
+ width (`int`):
+ Width of the input image.
+ images_kwargs (`dict`, *optional*)
+ Any kwargs to override defaults of the image processor.
+ Returns:
+ `int`: Number of image patches per image.
+ """
+ min_pixels = images_kwargs["min_pixels"] if "min_pixels" in images_kwargs else self.size["shortest_edge"]
+ max_pixels = images_kwargs["max_pixels"] if "max_pixels" in images_kwargs else self.size["longest_edge"]
+ patch_size = images_kwargs.get("patch_size", self.patch_size)
+ merge_size = images_kwargs.get("merge_size", self.merge_size)
+
+ factor = patch_size * merge_size
+ resized_height, resized_width = smart_resize(
+ height, width, factor, min_pixels=min_pixels, max_pixels=max_pixels
+ )
+ grid_h, grid_w = resized_height // patch_size, resized_width // patch_size
+ return grid_h * grid_w
+
+
+__all__ = ["GlmImageImageProcessorPil"]
diff --git a/third_party/transformers/src/transformers/models/glm_image/modeling_glm_image.py b/third_party/transformers/src/transformers/models/glm_image/modeling_glm_image.py
new file mode 100644
index 0000000000000000000000000000000000000000..7215076ea8cf0a5886c55e561e37eeb76df782ca
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_image/modeling_glm_image.py
@@ -0,0 +1,1728 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/glm_image/modular_glm_image.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_glm_image.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub, use_kernelized_func
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
+from ...utils.generic import maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_glm_image import GlmImageConfig, GlmImageTextConfig, GlmImageVisionConfig, GlmImageVQVAEConfig
+
+
+class GlmImageVisionMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.activation_fn = ACT2FN[config.hidden_act]
+ self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+ self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ return hidden_states
+
+
+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)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class GlmImageVisionAttention(nn.Module):
+ def __init__(self, config: GlmImageVisionConfig) -> None:
+ super().__init__()
+ self.dim = config.hidden_size
+ self.num_heads = config.num_heads
+ self.head_dim = self.dim // self.num_heads
+ self.num_key_value_groups = 1 # needed for eager attention
+ self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.attention_bias)
+ self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)
+ self.scaling = self.head_dim**-0.5
+ self.config = config
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ **kwargs,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ query_states, key_states, value_states = (
+ self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+ )
+ query_states = query_states.transpose(0, 1).unsqueeze(0)
+ key_states = key_states.transpose(0, 1).unsqueeze(0)
+ value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ if "flash" in self.config._attn_implementation:
+ # Flash Attention: Use cu_seqlens for variable length attention
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+ attn_output, _ = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ cu_seq_lens_q=cu_seqlens,
+ cu_seq_lens_k=cu_seqlens,
+ max_length_q=max_seqlen,
+ max_length_k=max_seqlen,
+ is_causal=False,
+ **kwargs,
+ )
+ else:
+ # Other implementations: Process each chunk separately
+ lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+ splits = [
+ torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
+ ]
+
+ attn_outputs = [
+ attention_interface(
+ self,
+ q,
+ k,
+ v,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ is_causal=False,
+ **kwargs,
+ )[0]
+ for q, k, v in zip(*splits)
+ ]
+ attn_output = torch.cat(attn_outputs, dim=1)
+
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+class GlmImageVisionPatchEmbed(nn.Module):
+ def __init__(self, config: GlmImageVisionConfig) -> None:
+ super().__init__()
+ self.patch_size = config.patch_size
+ self.in_channels = config.in_channels
+ self.embed_dim = config.hidden_size
+ kernel_size = [self.patch_size, self.patch_size]
+ self.proj = nn.Conv2d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size)
+
+ def forward(self, hidden_states) -> torch.Tensor:
+ target_dtype = self.proj.weight.dtype
+ hidden_states = hidden_states.view(-1, self.in_channels, self.patch_size, self.patch_size)
+ hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
+ return hidden_states
+
+
+class GlmImageVisionEmbeddings(nn.Module):
+ def __init__(self, config: GlmImageVisionConfig) -> None:
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.image_size = config.image_size
+ self.patch_size = config.patch_size
+
+ self.num_patches = (self.image_size // self.patch_size) ** 2
+ self.num_positions = self.num_patches
+ self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+ self.interpolated_method = "bilinear"
+
+ def forward(self, embeddings, lengths, image_shapes, h_coords, w_coords) -> torch.Tensor:
+ """
+ Forward pass with integrated position encoding adaptation using 2D interpolation.
+
+ Args:
+ embeddings: Input embeddings tensor
+ lengths (torch.Tensor): Sequence lengths for each image in the batch.
+ image_shapes (torch.Tensor): Tensor of shape [batch_size, 3] representing the image shapes (t, h, w).
+ h_coords (torch.Tensor): Tensor of shape [total_seq] representing the h coordinate for each patch.
+ w_coords (torch.Tensor): Tensor of shape [total_seq] representing the w coordinate for each patch.
+
+ Returns:
+ torch.Tensor: Embeddings with adapted position encoding added.
+ """
+ # Get position embedding parameters
+ pos_embed_weight = self.position_embedding.weight
+ hidden_size = pos_embed_weight.shape[1]
+ device = pos_embed_weight.device
+
+ # Convert inputs to tensors if needed
+ if isinstance(lengths, list):
+ lengths = torch.tensor(lengths, device=device, dtype=torch.long)
+
+ # Prepare 2D position embedding
+ orig_size_sq = pos_embed_weight.shape[0]
+ orig_size = int(orig_size_sq**0.5)
+ pos_embed_2d = (
+ pos_embed_weight.view(orig_size, orig_size, hidden_size)
+ .permute(2, 0, 1)
+ .unsqueeze(0)
+ .to(device=device, dtype=torch.float32)
+ )
+
+ # Calculate target dimensions for each patch
+ target_h = torch.cat([image_shapes[i, 1].repeat(lengths[i]) for i in range(len(lengths))]).to(
+ device=device, dtype=torch.float32
+ )
+ target_w = torch.cat([image_shapes[i, 2].repeat(lengths[i]) for i in range(len(lengths))]).to(
+ device=device, dtype=torch.float32
+ )
+
+ # Normalize coordinates to [-1, 1] range for grid_sample
+ norm_w = ((w_coords + 0.5) / target_w) * 2 - 1
+ norm_h = ((h_coords + 0.5) / target_h) * 2 - 1
+
+ # Create sampling grid
+ grid = torch.stack((norm_w, norm_h), dim=-1).unsqueeze(0).unsqueeze(2)
+
+ # Perform bicubic interpolation
+ interpolated_embed_fp32 = F.grid_sample(
+ pos_embed_2d, grid, mode=self.interpolated_method, align_corners=False, padding_mode="border"
+ )
+
+ # Reshape and convert back to original dtype
+ adapted_pos_embed_fp32 = interpolated_embed_fp32.squeeze(0).squeeze(-1).permute(1, 0)
+ adapted_pos_embed = adapted_pos_embed_fp32.to(pos_embed_weight.dtype).to(embeddings.device)
+
+ # Add adapted position encoding to embeddings
+ embeddings = embeddings + adapted_pos_embed
+ return embeddings
+
+
+class GlmImageVisionBlock(GradientCheckpointingLayer):
+ def __init__(self, config: GlmImageVisionConfig) -> None:
+ super().__init__()
+ self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.attn = GlmImageVisionAttention(config)
+ self.mlp = GlmImageVisionMLP(config)
+
+ @auto_docstring
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ r"""
+ cu_seqlens (`torch.Tensor` of shape `(num_images_or_videos + 1,)`):
+ The cumulative sequence lengths of each image or video feature.
+ """
+ residual = hidden_states
+
+ hidden_states = self.norm1(hidden_states)
+ hidden_states = self.attn(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.norm2(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+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)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+
+ # Keep half or full tensor for later concatenation
+ rotary_dim = cos.shape[-1]
+ q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+ k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+ # Apply rotary embeddings on the first half or full tensor
+ q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+ k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+
+ # Concatenate back to full shape
+ q_embed = torch.cat([q_embed, q_pass], dim=-1)
+ k_embed = torch.cat([k_embed, k_pass], dim=-1)
+ return q_embed, k_embed
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class GlmImageTextAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: GlmImageTextConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = True
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
+ self.rope_parameters = config.rope_parameters
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: torch.Tensor | None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape)
+ key_states = self.k_proj(hidden_states).view(hidden_shape)
+ value_states = self.v_proj(hidden_states).view(hidden_shape)
+
+ query_states = query_states.transpose(1, 2)
+ key_states = key_states.transpose(1, 2)
+ value_states = value_states.transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+@auto_docstring
+class GlmImagePreTrainedModel(PreTrainedModel):
+ config: GlmImageConfig
+ base_model_prefix = "model"
+ input_modalities = ("image", "text")
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["GlmImageTextDecoderLayer", "GlmImageVisionBlock"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn = True
+ _supports_sdpa = True
+
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ super()._init_weights(module)
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for Llava outputs, with hidden states and attentions.
+ """
+)
+class GlmImageModelOutputWithPast(ModelOutput):
+ r"""
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ last_hidden_state: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ rope_deltas: torch.LongTensor | None = None
+
+
+class GlmImageVQVAEVectorQuantizer(nn.Module):
+ """
+ A module for vector quantization using learned embedding vectors.
+
+ This module implements the quantization process similar to te one described in
+ the VQ-VAE (Vector Quantized Variational AutoEncoder) paper. It quantizes continuous
+ input vectors into discrete codebook vectors, which are learned during training.
+ Current implementation improves over previous ones by avoiding costly matrix multiplications
+ and allowing for post-hoc remapping of indices.
+ """
+
+ def __init__(self, config: GlmImageVQVAEConfig):
+ super().__init__()
+ self.num_embeddings = config.num_embeddings
+ self.embedding_dim = config.embed_dim
+ self.beta = getattr(config, "beta", 0.25)
+
+ self.embedding = nn.Embedding(self.num_embeddings, self.embedding_dim)
+
+ def forward(self, hidden_state: torch.Tensor):
+ hidden_state = hidden_state.permute(0, 2, 3, 1).contiguous()
+ hidden_state_flattened = hidden_state.view(-1, self.embedding_dim)
+
+ # L2 normalize
+ hidden_state = F.normalize(hidden_state, p=2, dim=-1)
+ hidden_state_flattened = F.normalize(hidden_state_flattened, p=2, dim=-1)
+ embedding = F.normalize(self.embedding.weight, p=2, dim=-1)
+
+ # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+ distances = (
+ torch.sum(hidden_state_flattened**2, dim=1, keepdim=True)
+ + torch.sum(embedding**2, dim=1)
+ - 2 * torch.einsum("bd,dn->bn", hidden_state_flattened, embedding.transpose(0, 1))
+ )
+
+ min_encoding_indices = torch.argmin(distances, dim=1)
+ hidden_state_quant = embedding[min_encoding_indices].view(hidden_state.shape)
+
+ # compute loss for embedding
+ loss = torch.mean((hidden_state_quant.detach() - hidden_state) ** 2) + self.beta * torch.mean(
+ (hidden_state_quant - hidden_state.detach()) ** 2
+ )
+
+ # preserve gradients
+ hidden_state_quant = hidden_state + (hidden_state_quant - hidden_state).detach()
+
+ # reshape back to match original input shape
+ hidden_state_quant = hidden_state_quant.permute(0, 3, 1, 2).contiguous()
+
+ return hidden_state_quant, loss, min_encoding_indices
+
+
+@dataclass
+@auto_docstring
+class GlmImageVQVAEModelOutput(BaseModelOutputWithPooling):
+ r"""
+ quantized_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ Quantized last hidden state from the VQ-VAE model.
+ image_tokens (`torch.FloatTensor` of shape `(batch_size, config.vocab_size`):
+ Indices of the image tokens predicted by the VQ-VAE model.
+ embedding_loss (`torch.FloatTensor`):
+ The embedding loss computed during quantization.
+ """
+
+ quantized_last_hidden_state: torch.FloatTensor | None = None
+ image_tokens: torch.FloatTensor | None = None
+ embedding_loss: torch.FloatTensor | None = None
+
+
+@auto_docstring(
+ custom_intro="""
+ The VQ-VAE model used in GlmImage for encoding/decoding images into discrete tokens.
+ This model follows the "Make-a-scene: Scene-based text-to-image generation with human priors" paper from
+ [ Oran Gafni, Adam Polyak, Oron Ashual, Shelly Sheynin, Devi Parikh, and Yaniv
+ Taigman](https://huggingface.co/papers/2203.13131).
+ """
+)
+class GlmImageVQVAE(GlmImagePreTrainedModel):
+ config: GlmImageVQVAEConfig
+ _no_split_modules = [
+ "GlmImageVQVAEVectorQuantizer",
+ ]
+ _can_record_outputs = {}
+
+ def __init__(self, config: GlmImageVQVAEConfig):
+ super().__init__(config)
+ self.quantize = GlmImageVQVAEVectorQuantizer(config)
+ self.quant_conv = torch.nn.Conv2d(config.latent_channels, config.embed_dim, 1)
+ self.post_quant_conv = torch.nn.Conv2d(config.embed_dim, config.latent_channels, 1)
+ self.eval() # GlmImage's VQ model is frozen
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def encode(self, hidden_states) -> GlmImageVQVAEModelOutput:
+ conv_hidden_states = self.quant_conv(hidden_states)
+ quantized_last_hidden_state, emb_loss, indices = self.quantize(conv_hidden_states)
+ return GlmImageVQVAEModelOutput(
+ last_hidden_state=hidden_states,
+ quantized_last_hidden_state=quantized_last_hidden_state,
+ image_tokens=indices,
+ embedding_loss=emb_loss,
+ )
+
+
+class GlmImageVisionModel(GlmImagePreTrainedModel):
+ config: GlmImageVisionConfig
+ input_modalities = ("image",)
+ _no_split_modules = ["GlmImageVisionBlock"]
+ _can_record_outputs = {
+ "hidden_states": GlmImageVisionBlock,
+ "attentions": GlmImageVisionAttention,
+ }
+ main_input_name = "pixel_values"
+
+ def __init__(self, config: GlmImageVisionConfig) -> None:
+ super().__init__(config)
+ self.spatial_merge_size = config.spatial_merge_size
+ self.patch_size = config.patch_size
+
+ self.embeddings = GlmImageVisionEmbeddings(config)
+ self.patch_embed = GlmImageVisionPatchEmbed(config)
+
+ head_dim = config.hidden_size // config.num_heads
+
+ self.blocks = nn.ModuleList([GlmImageVisionBlock(config) for _ in range(config.depth)])
+
+ self.gradient_checkpointing = False
+ self.head_dim = head_dim
+ self.post_init()
+
+ def rot_pos_emb(self, grid_thw):
+ pos_ids = []
+ for t, h, w in grid_thw:
+ hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
+ hpos_ids = hpos_ids.reshape(
+ h // self.spatial_merge_size,
+ self.spatial_merge_size,
+ w // self.spatial_merge_size,
+ self.spatial_merge_size,
+ )
+ hpos_ids = hpos_ids.permute(0, 2, 1, 3)
+ hpos_ids = hpos_ids.flatten()
+
+ wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
+ wpos_ids = wpos_ids.reshape(
+ h // self.spatial_merge_size,
+ self.spatial_merge_size,
+ w // self.spatial_merge_size,
+ self.spatial_merge_size,
+ )
+ wpos_ids = wpos_ids.permute(0, 2, 1, 3)
+ wpos_ids = wpos_ids.flatten()
+ pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
+ pos_ids = torch.cat(pos_ids, dim=0)
+ return pos_ids
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.Tensor` of shape `(total_patches, num_channels * patch_size * patch_size)`):
+ Packed pixel values.
+ grid_thw (`torch.Tensor` of shape `(num_images, 3)`):
+ The temporal, height and width of feature shape of each image.
+
+ Returns:
+ `torch.Tensor` of shape `(total_patches, hidden_size)`: Hidden states.
+ """
+
+ hidden_states = self.patch_embed(pixel_values)
+ image_type_ids = self.rot_pos_emb(grid_thw)
+
+ cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+ dim=0,
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+ )
+ cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+ seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+ hidden_states = self.embeddings(
+ hidden_states,
+ seqlens,
+ grid_thw,
+ image_type_ids[:, 0].to(hidden_states.device),
+ image_type_ids[:, 1].to(hidden_states.device),
+ )
+
+ # Transformer blocks (no position_embeddings needed, already added above)
+ for blk in self.blocks:
+ hidden_states = blk(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ )
+
+ return BaseModelOutputWithPooling(last_hidden_state=hidden_states)
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class GlmImageRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ GlmImageRMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class GlmImageTextRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: GlmImageTextConfig, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+ self.mrope_section = config.rope_parameters.get("mrope_section", [8, 12, 12])
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: GlmImageTextConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
+ head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+ dim = int(head_dim * partial_rotary_factor)
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ # In contrast to other models, GLM-V has different position ids for the grids
+ # So we expand the inv_freq to shape (3, ...)
+ inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
+ position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions)
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
+ freqs = self.apply_mrope(freqs, self.mrope_section)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+ def apply_mrope(self, freqs, mrope_section):
+ section = mrope_section
+ chunks = freqs.split(section, dim=-1)
+ result = torch.cat([chunk[i % 3] for i, chunk in enumerate(chunks)], dim=-1)
+ return result
+
+
+class GlmImageTextMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ self.config = config
+ self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
+ self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+ self.activation_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+ up_states = self.gate_up_proj(hidden_states)
+
+ gate, up_states = up_states.chunk(2, dim=-1)
+ up_states = up_states * self.activation_fn(gate)
+
+ return self.down_proj(up_states)
+
+
+class GlmImageTextDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: GlmImageTextConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = GlmImageTextAttention(config, layer_idx)
+ self.mlp = GlmImageTextMLP(config)
+ self.input_layernorm = GlmImageRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = GlmImageRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_self_attn_layernorm = GlmImageRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_mlp_layernorm = GlmImageRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ @auto_docstring
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ **kwargs,
+ ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = self.post_self_attn_layernorm(hidden_states)
+ 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 = self.post_mlp_layernorm(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+@auto_docstring
+class GlmImageTextModel(GlmImagePreTrainedModel):
+ config: GlmImageTextConfig
+ input_modalities = ("text",)
+ _can_record_outputs = {
+ "hidden_states": GlmImageTextDecoderLayer,
+ "attentions": GlmImageTextAttention,
+ }
+
+ def __init__(self, config: GlmImageTextConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList(
+ [GlmImageTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = GlmImageRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = GlmImageTextRotaryEmbedding(config=config)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple | BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ # torch.jit.trace() doesn't support cache objects in the output
+ if use_cache and past_key_values is None and not torch.jit.is_tracing():
+ past_key_values = DynamicCache(config=self.config)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ # the hard coded `3` is for temporal, height and width.
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
+ elif position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+ # NOTE: we need to pass text position ids for packing. Qwen2-VL uses 3D positions
+ # where each dim indicates visual spatial positions for temporal/height/width grids.
+ # There are two scenarios when FA2-like packed masking might be activated.
+ # 1. User specifically passed packed `position_ids` and no attention mask.
+ # In this case we expect the useer to create correct position ids for all 3 grids
+ # and prepend text-only position ids to it. The final tensor will be [4, bs, seq-len]
+ # 2. User runs forward with no attention mask and no position ids. In this case, position ids
+ # are prepared by the model (`get_rope_index`) as `[4, bs, seq-len]` tensor. Text-only positions are
+ # prepended by us when creating positions so that the mask is constructed correctly. NOTE: failing to pass
+ # text-only positions will cause incorrect mask construction, do not change `prepare_input_for_generation`
+ if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+ text_position_ids = position_ids[0]
+ position_ids = position_ids[1:]
+ else:
+ # If inputs are not packed (usual 3D positions), do not prepare mask from position_ids
+ text_position_ids = None
+
+ mask_kwargs = {
+ "config": self.config,
+ "inputs_embeds": inputs_embeds,
+ "attention_mask": attention_mask,
+ "past_key_values": past_key_values,
+ "position_ids": text_position_ids,
+ }
+ # Create the masks
+ causal_mask = create_causal_mask(**mask_kwargs)
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for decoder_layer in self.layers:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_ids=text_position_ids,
+ past_key_values=past_key_values,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = layer_outputs
+
+ hidden_states = self.norm(hidden_states)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class GlmImageModel(GlmImagePreTrainedModel):
+ base_model_prefix = "model"
+ # Reference: fix gemma3 grad acc #37208
+ accepts_loss_kwargs = False
+ _no_split_modules = ["GlmImageTextDecoderLayer", "GlmImageVisionBlock"]
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.visual = GlmImageVisionModel._from_config(config.vision_config)
+ self.language_model = GlmImageTextModel._from_config(config.text_config)
+
+ self.rope_deltas = None # cache rope_deltas here
+ self.vqmodel = GlmImageVQVAE._from_config(config.vq_config)
+
+ # Per-sample caches for batch processing
+ self._cached_decode_position_ids = None # shape: [batch_size, 3, max_decode_len]
+ self._prefill_len = None # prefill sequence length (same for all samples in batch)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.language_model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.language_model.set_input_embeddings(value)
+
+ def get_vision_position_ids(
+ self,
+ start_position: int,
+ grid_thw: list[int, int, int] | torch.Tensor,
+ temp_merge_size: int = 1,
+ spatial_merge_size: int = 1,
+ time_interval: int = 1,
+ device: str | torch.device | None = None,
+ ):
+ """
+ Compute 3D positional indices for vision tokens derived from a single image or video input.
+
+ The positions are generated from the input grid defined by temporal (T), height (H), and
+ width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the
+ merge sizes used in the vision backbone. The resulting positions are offset by `start_position`.
+
+ Args:
+ start_position (`int`):
+ Offset added to all computed positional indices.
+ grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`):
+ The (T, H, W) grid representing the feature layout of the current image or video after patch embedding.
+ temp_merge_size (`int`, *optional*):
+ Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided
+ by this value. Defaults to 1.
+ spatial_merge_size (`int`, *optional*):
+ Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided
+ by this value. Defaults to 1.
+ time_interval (`int`, *optional*):
+ Spacing factor applied between consecutive temporal position indices.Defaults to 1.
+ device (`str` or `torch.device`, *optional*):
+ Device on which the resulting tensor is allocated. If `None`, uses the current default device.
+
+ Returns:
+ torch.LongTensor of shape (3, sequence_length):
+ Positional indices for temporal, height, and width dimensions,
+ flattened into sequence form and offset by `start_position`.
+ """
+ llm_grid_t, llm_grid_h, llm_grid_w = (
+ grid_thw[0].item() // temp_merge_size,
+ grid_thw[1].item() // spatial_merge_size,
+ grid_thw[2].item() // spatial_merge_size,
+ )
+
+ image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t
+ position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat(
+ llm_grid_h * llm_grid_t
+ )
+ position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave(
+ llm_grid_w * llm_grid_t
+ )
+ position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long)
+ position_temporal = position_temporal * time_interval
+ vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0)
+
+ return vision_position_ids
+
+ def get_rope_index(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ images_per_sample: torch.LongTensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Calculate the 3D rope index for image generation task with full batch support.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary.
+ image_grid_thw (`torch.LongTensor` of shape `(total_images_in_batch, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image.
+ Images are packed across all samples in the batch.
+ images_per_sample (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Number of images (including target grids) for each sample in the batch.
+ Used to split image_grid_thw by sample.
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding token indices.
+
+ Returns:
+ position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`):
+ Position IDs for temporal, height, and width dimensions.
+ mrope_position_deltas (`torch.Tensor` of shape `(batch_size, 1)`):
+ Position deltas for multi-modal rotary position embedding.
+ """
+ batch_size, seq_len = input_ids.shape
+ device = input_ids.device
+ dtype = input_ids.dtype
+
+ image_start_token_id = self.config.image_start_token_id
+ image_end_token_id = self.config.image_end_token_id
+
+ position_ids = torch.ones(3, batch_size, seq_len, dtype=dtype, device=device)
+ text_positions = torch.arange(seq_len, device=device)[None, :].repeat(3, 1)
+
+ # Split image_grid_thw by sample if images_per_sample is provided
+ if image_grid_thw is not None and images_per_sample is not None:
+ grids_per_sample = torch.split(image_grid_thw, images_per_sample.tolist())
+ elif image_grid_thw is not None:
+ # Fallback: assume all grids belong to first sample (batch_size=1)
+ grids_per_sample = [image_grid_thw] * batch_size
+ else:
+ grids_per_sample = [None] * batch_size
+
+ # Per-sample caches for decode stage
+ all_decode_position_ids = []
+
+ for batch_idx in range(batch_size):
+ curr_input_ids = input_ids[batch_idx]
+ curr_grids = grids_per_sample[batch_idx]
+
+ if attention_mask is not None and attention_mask.shape[1] == seq_len:
+ valid_mask = attention_mask[batch_idx] == 1
+ curr_input_ids_valid = curr_input_ids[valid_mask]
+ else:
+ # attention_mask may have different length during assisted decoding
+ curr_input_ids_valid = curr_input_ids
+ valid_mask = None
+
+ # Find image boundaries in this sample
+ image_end_positions = torch.where(curr_input_ids_valid == image_end_token_id)[0]
+ image_start_positions = torch.where(curr_input_ids_valid == image_start_token_id)[0] + 1
+ num_complete_images = len(image_end_positions)
+
+ current_pos = 0
+ prev_image_end = 0
+ curr_position_ids = []
+
+ # Process complete images (source images in image-to-image task)
+ for img_idx, (start, end) in enumerate(zip(image_start_positions, image_end_positions)):
+ if curr_grids is None or img_idx >= len(curr_grids):
+ break
+
+ # Text tokens before this image
+ llm_pos_length = start - prev_image_end
+ llm_position_ids = text_positions[:, current_pos : current_pos + llm_pos_length].to(device=device)
+ current_pos += llm_position_ids.shape[-1]
+
+ # Image tokens with 2D spatial encoding
+ # For an image with height H and width W:
+ # - position_width cycles [0, 1, ..., W-1] for each row, repeated H times
+ # - position_height stays constant per row, [0]*W, [1]*W, ..., [H-1]*W
+ vision_position_ids = self.get_vision_position_ids(
+ start_position=current_pos, grid_thw=curr_grids[img_idx], device=device
+ )
+ current_pos += max(curr_grids[img_idx][1], curr_grids[img_idx][2])
+
+ prev_image_end = end
+ curr_position_ids.append(torch.cat([llm_position_ids, vision_position_ids], dim=-1))
+
+ # Remaining text tokens (including the final image_start token for generation)
+ end_position = len(curr_input_ids_valid) - prev_image_end
+ llm_position_ids = text_positions[:, current_pos : current_pos + end_position].to(device=device)
+ current_pos += llm_position_ids.shape[-1]
+ curr_position_ids.append(llm_position_ids)
+
+ # Concatenate all position ids for this sample
+ curr_position_ids = torch.cat(curr_position_ids, dim=-1)
+
+ # Store in the main position_ids tensor
+ if valid_mask is not None:
+ position_ids[:, batch_idx, valid_mask] = curr_position_ids
+ else:
+ position_ids[:, batch_idx, :] = curr_position_ids
+
+ # Build decode position ids for this sample
+ if curr_grids is not None and len(curr_grids) > 0:
+ num_decode_grids = len(curr_grids) - num_complete_images
+ num_decode_grids = max(num_decode_grids, 0)
+ decode_pos = current_pos
+
+ decode_temporal_list = []
+ decode_height_list = []
+ decode_width_list = []
+
+ curr_grids_list = curr_grids.tolist()
+ for i in range(1, num_decode_grids + 1):
+ grid_idx = -i
+ h = curr_grids_list[grid_idx][1]
+ w = curr_grids_list[grid_idx][2]
+ total_tokens = h * w
+
+ h_indices = torch.arange(h, device=device).unsqueeze(1).expand(h, w).flatten()
+ w_indices = torch.arange(w, device=device).unsqueeze(0).expand(h, w).flatten()
+
+ decode_temporal_list.append(
+ torch.full((total_tokens,), decode_pos, device=device, dtype=torch.long)
+ )
+ decode_height_list.append(decode_pos + h_indices)
+ decode_width_list.append(decode_pos + w_indices)
+ decode_pos = decode_pos + max(h, w)
+
+ # End marker
+ decode_temporal_list.append(torch.tensor([decode_pos], device=device, dtype=torch.long))
+ decode_height_list.append(torch.tensor([decode_pos], device=device, dtype=torch.long))
+ decode_width_list.append(torch.tensor([decode_pos], device=device, dtype=torch.long))
+
+ sample_decode_pos_ids = torch.stack(
+ [
+ torch.cat(decode_temporal_list, dim=0),
+ torch.cat(decode_height_list, dim=0),
+ torch.cat(decode_width_list, dim=0),
+ ],
+ dim=0,
+ )
+ all_decode_position_ids.append(sample_decode_pos_ids)
+
+ # Store prefill length (same for all samples since input_ids is padded to same length)
+ self._prefill_len = seq_len
+
+ # Pad decode position ids to same length and stack
+ if all_decode_position_ids:
+ max_decode_len = max(x.shape[1] for x in all_decode_position_ids)
+ padded_decode_pos_ids = [
+ F.pad(pos_ids, (0, max_decode_len - pos_ids.shape[1]), mode="replicate")
+ for pos_ids in all_decode_position_ids
+ ]
+ self._cached_decode_position_ids = torch.stack(padded_decode_pos_ids, dim=0) # [batch, 3, max_decode_len]
+ else:
+ self._cached_decode_position_ids = None
+
+ mrope_position_deltas = torch.zeros([batch_size, 1], dtype=dtype, device=device)
+
+ return position_ids, mrope_position_deltas
+
+ @can_return_tuple
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ pixel_values = pixel_values.type(self.visual.dtype)
+ vision_outputs = self.visual(pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs)
+ split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
+ image_embeds = torch.split(vision_outputs.last_hidden_state, split_sizes)
+ vision_outputs.pooler_output = image_embeds
+
+ return vision_outputs
+
+ def get_placeholder_mask(
+ self,
+ input_ids: torch.LongTensor,
+ image_ids: torch.LongTensor,
+ ):
+ """
+ Replace image placeholder tokens in input_ids with actual image token ids from VQVAE.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, seq_len)`):
+ Input token ids with image placeholders.
+ image_ids (`torch.LongTensor` of shape `(num_images, num_tokens_per_image)` or flattened):
+ Discrete token indices from the VQVAE codebook.
+
+ Returns:
+ special_image_mask (`torch.LongTensor` of shape `(batch_size, seq_len)`):
+ Mask indicating positions in input ids that will be replaced by actual image tokens.
+ """
+
+ special_image_mask = input_ids == self.config.image_token_id
+ n_placeholder_tokens = special_image_mask.sum()
+ n_image_tokens = image_ids.shape[0]
+
+ if n_placeholder_tokens != n_image_tokens:
+ raise ValueError(
+ f"Number of image placeholder tokens ({n_placeholder_tokens.item()}) does not match "
+ f"number of image tokens from VQVAE ({n_image_tokens})"
+ )
+
+ return special_image_mask
+
+ def compute_3d_position_ids(
+ self,
+ input_ids: torch.Tensor | None,
+ image_grid_thw: torch.Tensor | None,
+ images_per_sample: torch.Tensor | None,
+ inputs_embeds: torch.Tensor | None,
+ attention_mask: torch.Tensor | None,
+ past_key_values: torch.Tensor | None,
+ ) -> torch.Tensor | None:
+ past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length()
+ can_compute_mrope = input_ids is not None and image_grid_thw is not None
+
+ if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0):
+ position_ids, rope_deltas = self.get_rope_index(
+ input_ids,
+ image_grid_thw=image_grid_thw,
+ attention_mask=attention_mask,
+ images_per_sample=images_per_sample,
+ )
+ self.rope_deltas = rope_deltas
+ # Use pre-calculated rope-deltas to infer correct 3D position ids during incremental
+ # generation (past_key_values_length > 0) or when only inputs_embeds is provided (no input_ids
+ # to recompute from). Skip when input_ids is provided without past_key_values to avoid shape
+ # mismatches from stale rope_deltas (e.g., training forward pass after generation).
+ elif self.rope_deltas is not None and (past_key_values_length > 0 or input_ids is None):
+ batch_size, seq_length, _ = inputs_embeds.shape
+ if self._cached_decode_position_ids is not None:
+ step = past_key_values_length - self._prefill_len
+ position_ids = self._cached_decode_position_ids[:, :, step : step + seq_length].permute(1, 0, 2)
+ else:
+ position_ids = (
+ torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_key_values_length
+ )
+ position_ids = position_ids.view(1, 1, -1).repeat(3, batch_size, 1)
+ else:
+ # Can't build correct 3D positions. Let the model infer it
+ position_ids = None
+ return position_ids
+
+ @auto_docstring
+ @can_return_tuple
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ pixel_values: torch.Tensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ images_per_sample: torch.LongTensor | None = None,
+ rope_deltas: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | GlmImageModelOutputWithPast:
+ r"""
+ image_grid_thw (`torch.LongTensor` of shape `(total_images_in_batch, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ Images are packed across all samples in the batch.
+ images_per_sample (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Number of images (including target grids) for each sample in the batch.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ batch_size = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
+
+ if pixel_values is not None:
+ # Process source images (image-to-image mode)
+ # Source images are identified by counting image_end_token_id in input_ids
+ # Note: We must exclude padding tokens since pad_token_id == image_end_token_id
+ if images_per_sample is not None:
+ grids_per_sample = torch.split(image_grid_thw, images_per_sample.tolist())
+ # Create mask for non-padding tokens (attention_mask=1 means non-padding)
+ # Handle 4D attention mask (from static cache) by extracting diagonal
+ if attention_mask is not None and attention_mask.ndim == 4:
+ non_pad_mask = torch.diagonal(attention_mask[:, 0], dim1=1, dim2=2)
+ if non_pad_mask.dtype.is_floating_point:
+ non_pad_mask = non_pad_mask / torch.finfo(non_pad_mask.dtype).min
+ non_pad_mask = (1.0 - non_pad_mask).int()
+ # Only keep columns matching input_ids length
+ non_pad_mask = non_pad_mask[:, -input_ids.shape[1] :]
+ else:
+ non_pad_mask = attention_mask if attention_mask is not None else torch.ones_like(input_ids)
+
+ source_grids_list = []
+ is_image_end = input_ids == self.config.image_end_token_id
+ is_non_pad = non_pad_mask == 1
+ num_source_per_sample = (is_image_end & is_non_pad).sum(dim=1).tolist()
+ for sample_idx in range(batch_size):
+ num_source = num_source_per_sample[sample_idx]
+ if num_source > 0:
+ source_grids_list.append(grids_per_sample[sample_idx][:num_source])
+ if len(source_grids_list) == 0:
+ raise ValueError(
+ "pixel_values provided but no source images found in input_ids. "
+ "Ensure input_ids contains image_end_token_id for each source image."
+ )
+ source_grids = torch.cat(source_grids_list, dim=0)
+ else:
+ # Fallback for batch_size=1: all but last grid are source images
+ source_grids = image_grid_thw[:-1]
+
+ image_features = self.get_image_features(pixel_values, source_grids, return_dict=True)
+ image_embeds = torch.cat(image_features.pooler_output, dim=0)
+ image_ids = self.get_image_tokens(image_embeds, source_grids)
+ image_ids = image_ids.view(-1).to(input_ids.device)
+ special_image_mask = self.get_placeholder_mask(input_ids, image_ids)
+ input_ids = input_ids.masked_scatter(special_image_mask, image_ids)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if position_ids is None:
+ position_ids = self.compute_3d_position_ids(
+ input_ids=input_ids,
+ image_grid_thw=image_grid_thw,
+ images_per_sample=images_per_sample,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ )
+
+ outputs = self.language_model(
+ input_ids=None,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ return GlmImageModelOutputWithPast(
+ last_hidden_state=outputs.last_hidden_state,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ rope_deltas=self.rope_deltas,
+ )
+
+ def get_image_tokens(
+ self,
+ hidden_states: torch.FloatTensor,
+ image_grid_thw: torch.LongTensor,
+ ) -> torch.LongTensor:
+ """
+ Tokenizes image features into discrete tokens with VQVAE module.
+
+ Args:
+ hidden_states (`torch.FloatTensor` of shape `(total_patches, hidden_size)`):
+ The packed image features from vision encoder.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`):
+ The temporal, height and width of feature shape of each image.
+
+ Returns:
+ image_tokens (`torch.LongTensor` of shape `(total_patches,)`):
+ Discrete token indices from the VQVAE codebook.
+ """
+ hidden_size = hidden_states.shape[-1]
+ split_sizes = (image_grid_thw.prod(dim=-1)).tolist()
+ hidden_states_list = torch.split(hidden_states, split_sizes, dim=0)
+
+ all_image_toks = []
+ for i, hs in enumerate(hidden_states_list):
+ grid_t, grid_h, grid_w = image_grid_thw[i].tolist()
+ hs = hs.view(grid_t, grid_h, grid_w, hidden_size)
+ hs = hs.permute(0, 3, 1, 2).contiguous()
+ vqmodel_outputs: GlmImageVQVAEModelOutput = self.vqmodel.encode(hs)
+ all_image_toks.append(vqmodel_outputs.image_tokens)
+ return torch.cat(all_image_toks, dim=0)
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for GlmImage causal language model (or autoregressive) outputs.
+ """
+)
+class GlmImageCausalLMOutputWithPast(ModelOutput):
+ r"""
+ 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 (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ rope_deltas: torch.LongTensor | None = None
+
+
+class GlmImageForConditionalGeneration(GlmImagePreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {}
+ # Reference: fix gemma3 grad acc #37208
+ accepts_loss_kwargs = False
+ base_model_prefix = "model"
+ config: GlmImageConfig
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = GlmImageModel(config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vision_vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ return self.model.get_image_features(pixel_values, image_grid_thw, **kwargs)
+
+ def get_image_tokens(self, hidden_states: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None):
+ return self.model.get_image_tokens(hidden_states, image_grid_thw)
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ pixel_values: torch.Tensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ images_per_sample: torch.LongTensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | GlmImageCausalLMOutputWithPast:
+ r"""
+ 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]`.
+ image_grid_thw (`torch.LongTensor` of shape `(total_images_in_batch, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ Images are packed across all samples in the batch.
+ images_per_sample (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Number of images (including target grids) for each sample in the batch.
+
+ Example:
+
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoProcessor, GlmImageForConditionalGeneration
+
+ >>> model = GlmImageForConditionalGeneration.from_pretrained("zai-org/GLM-Image")
+ >>> processor = AutoProcessor.from_pretrained("zai-org/GLM-Image")
+
+ >>> messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "image"},
+ {"type": "text", "text": "Add a truck of this photo.28 40"},
+ ],
+ },
+ ]
+ >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+ >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos])
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."
+ ```"""
+ outputs = self.model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ image_grid_thw=image_grid_thw,
+ images_per_sample=images_per_sample,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)
+
+ return GlmImageCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ rope_deltas=outputs.rope_deltas,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ image_grid_thw=None,
+ images_per_sample=None,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ position_ids=position_ids,
+ pixel_values=pixel_values,
+ image_grid_thw=image_grid_thw,
+ is_first_iteration=is_first_iteration,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ model_inputs["position_ids"] = None
+ model_inputs["images_per_sample"] = images_per_sample
+
+ if not is_first_iteration and use_cache:
+ model_inputs["pixel_values"] = None
+
+ return model_inputs
+
+ def _get_image_nums(
+ self,
+ input_ids: torch.LongTensor | None,
+ ) -> torch.Tensor:
+ """
+ Get the number of images for each sample.
+ For GLM-Image, only input_ids allow us to get the number of images.
+
+ Returns:
+ image_counts (`torch.LongTensor` of shape `(batch_size,)`)
+ """
+ is_image = input_ids == self.config.image_start_token_id
+
+ return is_image.sum(dim=1)
+
+ def _expand_inputs_for_generation(
+ self,
+ expand_size: int = 1,
+ is_encoder_decoder: bool = False,
+ input_ids: torch.LongTensor | None = None,
+ **model_kwargs,
+ ) -> tuple[torch.LongTensor, dict[str, Any]]:
+ # Overwritten -- Support for expanding tensors without a batch size dimension
+ # e.g., pixel_values, image_grid_thw
+ # pixel_values.shape[0] is sum(seqlen_images for samples)
+ # image_grid_thw.shape[0] is sum(num_images for samples)
+
+ if expand_size == 1:
+ return input_ids, model_kwargs
+
+ visual_keys = ["pixel_values", "image_grid_thw", "images_per_sample"]
+
+ def _expand_dict_for_generation_visual(dict_to_expand):
+ image_grid_thw = model_kwargs.get("image_grid_thw", None)
+ if image_grid_thw is None:
+ return dict_to_expand
+
+ images_per_sample = model_kwargs.get("images_per_sample", None)
+
+ # Use images_per_sample if available
+ if images_per_sample is not None:
+ image_nums = images_per_sample.tolist()
+ elif input_ids is not None:
+ # Try to infer from image_grid_thw / batch_size
+ batch_size = input_ids.shape[0]
+ total_grids = image_grid_thw.shape[0]
+ if total_grids % batch_size == 0:
+ grids_per_sample = total_grids // batch_size
+ image_nums = [grids_per_sample] * batch_size
+ else:
+ # Cannot evenly distribute grids - fall back to simple repeat_interleave
+ # This handles test cases where image_grid_thw has (batch_size + 1) rows
+ dict_to_expand["image_grid_thw"] = image_grid_thw.repeat_interleave(expand_size, dim=0)
+ if dict_to_expand.get("pixel_values") is not None:
+ dict_to_expand["pixel_values"] = dict_to_expand["pixel_values"].repeat_interleave(
+ expand_size, dim=0
+ )
+ return dict_to_expand
+ else:
+ image_nums = self._get_image_nums(input_ids).tolist()
+
+ # Get source image counts per sample from image_end_token_id count
+ source_image_nums = (input_ids == self.config.image_end_token_id).sum(dim=1).tolist()
+
+ def _repeat_interleave_samples(x, lengths, repeat_times):
+ samples = torch.split(x, lengths)
+ repeat_args = [repeat_times] + [1] * (x.dim() - 1)
+ result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0)
+ return result
+
+ for key in dict_to_expand:
+ if key == "pixel_values":
+ # Split images into samples based on source image counts
+ if sum(source_image_nums) > 0:
+ # Split grids by sample to compute pixel counts
+ grids_per_sample = torch.split(image_grid_thw, image_nums)
+ all_pixel_counts = image_grid_thw.prod(dim=1)
+ pixel_counts_per_sample = torch.split(all_pixel_counts, image_nums)
+ # Build source mask and compute per-sample source pixel counts in one sync
+ source_pixel_counts = torch.zeros(len(grids_per_sample), device=image_grid_thw.device)
+ for batch_idx in range(len(grids_per_sample)):
+ num_source = source_image_nums[batch_idx]
+ if num_source > 0:
+ source_pixel_counts[batch_idx] = pixel_counts_per_sample[batch_idx][:num_source].sum()
+ lengths = source_pixel_counts.to(torch.int64).tolist()
+
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "image_grid_thw":
+ # Expand all grids (source + target) per sample
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=image_nums, repeat_times=expand_size
+ )
+ elif key == "images_per_sample":
+ # Simply repeat the counts
+ if dict_to_expand.get(key) is not None:
+ dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
+ return dict_to_expand
+
+ def _expand_dict_for_generation(dict_to_expand):
+ for key in dict_to_expand:
+ if (
+ dict_to_expand[key] is not None
+ and isinstance(dict_to_expand[key], torch.Tensor)
+ and key not in visual_keys
+ ):
+ dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
+ return dict_to_expand
+
+ model_kwargs = _expand_dict_for_generation_visual(model_kwargs)
+
+ if input_ids is not None:
+ input_ids = input_ids.repeat_interleave(expand_size, dim=0)
+
+ model_kwargs = _expand_dict_for_generation(model_kwargs)
+
+ if is_encoder_decoder:
+ if model_kwargs.get("encoder_outputs") is None:
+ raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.")
+ model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])
+
+ return input_ids, model_kwargs
+
+
+__all__ = [
+ "GlmImagePreTrainedModel",
+ "GlmImageVQVAE",
+ "GlmImageVisionModel",
+ "GlmImageTextModel",
+ "GlmImageModel",
+ "GlmImageForConditionalGeneration",
+]
diff --git a/third_party/transformers/src/transformers/models/glm_image/modular_glm_image.py b/third_party/transformers/src/transformers/models/glm_image/modular_glm_image.py
new file mode 100644
index 0000000000000000000000000000000000000000..e72aede3da661727919577822d0a62ddec7a94eb
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_image/modular_glm_image.py
@@ -0,0 +1,1492 @@
+# Copyright 2025 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+import math
+from collections.abc import Callable
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub.dataclasses import strict
+
+from ...cache_utils import Cache
+from ...configuration_utils import PreTrainedConfig
+from ...feature_extraction_utils import BatchFeature
+from ...generation import GenerationMixin
+from ...image_utils import ImageInput
+from ...modeling_outputs import BaseModelOutputWithPooling
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import ImagesKwargs, ProcessorMixin, Unpack
+from ...tokenization_utils_base import PreTokenizedInput, TextInput
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils.generic import merge_with_config_defaults
+from ...utils.import_utils import requires
+from ...utils.output_capturing import capture_outputs
+from ..chameleon.modeling_chameleon import ChameleonVQVAE, ChameleonVQVAEModelOutput, ChameleonVQVAEVectorQuantizer
+from ..glm4v.configuration_glm4v import Glm4vTextConfig, Glm4vVisionConfig
+from ..glm4v.modeling_glm4v import (
+ Glm4vCausalLMOutputWithPast,
+ Glm4vModel,
+ Glm4vModelOutputWithPast,
+ Glm4vPreTrainedModel,
+ Glm4vTextModel,
+ Glm4vVisionAttention,
+ Glm4vVisionBlock,
+ Glm4vVisionEmbeddings,
+ Glm4vVisionModel,
+ Glm4vVisionPatchEmbed,
+)
+from ..glm4v_moe.modeling_glm4v_moe import Glm4vMoeTextAttention, eager_attention_forward
+from ..qwen2_vl.image_processing_pil_qwen2_vl import Qwen2VLImageProcessorPil
+from ..qwen2_vl.image_processing_qwen2_vl import Qwen2VLImageProcessor
+from ..qwen2_vl.processing_qwen2_vl import Qwen2VLProcessorKwargs
+from ..siglip.modeling_siglip import SiglipMLP
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="zai-org/GLM-Image")
+@strict
+class GlmImageVQVAEConfig(PreTrainedConfig):
+ model_type = "glm_image_vqmodel"
+ base_config_key = "vq_config"
+
+ embed_dim: int = 2048
+ num_embeddings: int = 16384
+ latent_channels: int = 1536
+ in_channels: int = 3
+ initializer_range: float = 0.02
+
+
+@auto_docstring(checkpoint="zai-org/GLM-Image")
+@strict
+class GlmImageVisionConfig(Glm4vVisionConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import GlmImageVisionConfig, GlmImageVisionModel
+
+ >>> # Initializing a GlmImageVisionConfig GLM-Image style configuration
+ >>> configuration = GlmImageVisionConfig()
+
+ >>> # Initializing a model (with random weights) from the GLM-Image configuration
+ >>> model = GlmImageVisionModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "glm_image_vision"
+ base_config_key = "vision_config"
+
+ depth: int = 40
+ hidden_act: str = "gelu"
+ attention_bias: bool = True
+ num_heads: int = 16
+ image_size: int | list[int] | tuple[int, int] = 2048
+ patch_size: int | list[int] | tuple[int, int] = 16
+ layer_norm_eps: float = 1e-06
+ spatial_merge_size: int = 1
+ intermediate_size: int = 6144
+
+ out_hidden_size = AttributeError()
+ rms_norm_eps = AttributeError()
+ temporal_patch_size = AttributeError()
+
+
+@auto_docstring(checkpoint="zai-org/GLM-Image")
+@strict
+class GlmImageTextConfig(Glm4vTextConfig):
+ r"""
+ vision_vocab_size (`int`, *optional*, defaults to 16512):
+ Vision vocabulary size of the GlmImage model. Defines the number of different tokens that can be
+ represented by the `inputs_ids` passed when calling [`GlmImageVisionModel`]
+
+ Example:
+
+ ```python
+ >>> from transformers import GlmImageTextModel, GlmImageConfig
+
+ >>> # Initializing a GlmImageConfig style configuration
+ >>> configuration = GlmImageConfig()
+
+ >>> # Initializing a model from the GlmImageConfig style configuration
+ >>> model = GlmImageTextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ vocab_size: int = 168064
+ max_position_embeddings: int = 131072
+ vision_vocab_size: int = 16512
+ attention_bias: bool = True
+ pad_token_id: int = 167841
+ eos_token_id: int | list[int] | None = 16385
+
+
+@auto_docstring(checkpoint="zai-org/GLM-Image")
+@strict
+class GlmImageConfig(PreTrainedConfig):
+ r"""
+ image_start_token_id (`int`, *optional*, defaults to 16384):
+ The image start token index to encode the start of image.
+ image_end_token_id (`int`, *optional*, defaults to 16385):
+ The image end token index to encode the end of image.
+
+ ```python
+ >>> from transformers import Glm4vForConditionalGeneration, Glm4vConfig
+
+ >>> # Initializing a GLM-Image style configuration
+ >>> configuration = Glm4vConfig()
+
+ >>> # Initializing a model from the GLM-Image style configuration
+ >>> model = Glm4vForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "glm_image"
+ sub_configs = {
+ "vision_config": GlmImageVisionConfig,
+ "text_config": GlmImageTextConfig,
+ "vq_config": GlmImageVQVAEConfig,
+ }
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ text_config: dict | PreTrainedConfig | None = None
+ vision_config: dict | PreTrainedConfig | None = None
+ vq_config: dict | PreTrainedConfig | None = None
+ image_token_id: int = 167855
+ image_start_token_id: int = 16384
+ image_end_token_id: int = 16385
+ tie_word_embeddings: bool = False
+
+ def __post_init__(self, **kwargs):
+ if isinstance(self.vision_config, dict):
+ self.vision_config = self.sub_configs["vision_config"](**self.vision_config)
+ elif self.vision_config is None:
+ self.vision_config = self.sub_configs["vision_config"](**kwargs)
+
+ if isinstance(self.vq_config, dict):
+ self.vq_config = self.sub_configs["vq_config"](**self.vq_config)
+ elif self.vq_config is None:
+ self.vq_config = self.sub_configs["vq_config"](**kwargs)
+
+ if isinstance(self.text_config, dict):
+ self.text_config = self.sub_configs["text_config"](**self.text_config)
+ elif self.text_config is None:
+ self.text_config = self.sub_configs["text_config"](**kwargs)
+
+ super().__post_init__(**kwargs)
+
+
+class GlmImageVisionMLP(SiglipMLP):
+ pass
+
+
+class GlmImageVisionAttention(Glm4vVisionAttention):
+ def __init__(self, config: GlmImageVisionConfig) -> None:
+ super().__init__(config)
+ self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.attention_bias)
+ self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ **kwargs,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ query_states, key_states, value_states = (
+ self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+ )
+ query_states = query_states.transpose(0, 1).unsqueeze(0)
+ key_states = key_states.transpose(0, 1).unsqueeze(0)
+ value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ if "flash" in self.config._attn_implementation:
+ # Flash Attention: Use cu_seqlens for variable length attention
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+ attn_output, _ = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ cu_seq_lens_q=cu_seqlens,
+ cu_seq_lens_k=cu_seqlens,
+ max_length_q=max_seqlen,
+ max_length_k=max_seqlen,
+ is_causal=False,
+ **kwargs,
+ )
+ else:
+ # Other implementations: Process each chunk separately
+ lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+ splits = [
+ torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
+ ]
+
+ attn_outputs = [
+ attention_interface(
+ self,
+ q,
+ k,
+ v,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ is_causal=False,
+ **kwargs,
+ )[0]
+ for q, k, v in zip(*splits)
+ ]
+ attn_output = torch.cat(attn_outputs, dim=1)
+
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+class GlmImageVisionPatchEmbed(Glm4vVisionPatchEmbed):
+ def __init__(self, config: GlmImageVisionConfig) -> None:
+ super().__init__(config)
+
+ del self.temporal_patch_size
+ kernel_size = [self.patch_size, self.patch_size]
+ self.proj = nn.Conv2d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size)
+
+ def forward(self, hidden_states):
+ target_dtype = self.proj.weight.dtype
+ hidden_states = hidden_states.view(-1, self.in_channels, self.patch_size, self.patch_size)
+ hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
+ return hidden_states
+
+
+class GlmImageVisionEmbeddings(Glm4vVisionEmbeddings):
+ def __init__(self, config: GlmImageVisionConfig) -> None:
+ super().__init__(config)
+ self.interpolated_method = "bilinear"
+
+
+class GlmImageVisionBlock(Glm4vVisionBlock):
+ def __init__(self, config: GlmImageVisionConfig):
+ super().__init__(config)
+ self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.attn = GlmImageVisionAttention(config)
+ self.mlp = GlmImageVisionMLP(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ r"""
+ cu_seqlens (`torch.Tensor` of shape `(num_images_or_videos + 1,)`):
+ The cumulative sequence lengths of each image or video feature.
+ """
+ residual = hidden_states
+
+ hidden_states = self.norm1(hidden_states)
+ hidden_states = self.attn(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.norm2(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+class GlmImageTextAttention(Glm4vMoeTextAttention):
+ pass
+
+
+class GlmImagePreTrainedModel(Glm4vPreTrainedModel):
+ config: GlmImageConfig
+ input_modalities = ("image", "text")
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ PreTrainedModel._init_weights(module)
+
+
+class GlmImageModelOutputWithPast(Glm4vModelOutputWithPast):
+ pass
+
+
+class GlmImageVQVAEVectorQuantizer(ChameleonVQVAEVectorQuantizer):
+ def __init__(self, config: GlmImageVQVAEConfig):
+ super().__init__(config)
+ self.num_embeddings = config.num_embeddings
+ self.embedding_dim = config.embed_dim
+ self.beta = getattr(config, "beta", 0.25)
+
+ self.embedding = nn.Embedding(self.num_embeddings, self.embedding_dim)
+
+ def forward(self, hidden_state: torch.Tensor):
+ hidden_state = hidden_state.permute(0, 2, 3, 1).contiguous()
+ hidden_state_flattened = hidden_state.view(-1, self.embedding_dim)
+
+ # L2 normalize
+ hidden_state = F.normalize(hidden_state, p=2, dim=-1)
+ hidden_state_flattened = F.normalize(hidden_state_flattened, p=2, dim=-1)
+ embedding = F.normalize(self.embedding.weight, p=2, dim=-1)
+
+ # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z
+ distances = (
+ torch.sum(hidden_state_flattened**2, dim=1, keepdim=True)
+ + torch.sum(embedding**2, dim=1)
+ - 2 * torch.einsum("bd,dn->bn", hidden_state_flattened, embedding.transpose(0, 1))
+ )
+
+ min_encoding_indices = torch.argmin(distances, dim=1)
+ hidden_state_quant = embedding[min_encoding_indices].view(hidden_state.shape)
+
+ # compute loss for embedding
+ loss = torch.mean((hidden_state_quant.detach() - hidden_state) ** 2) + self.beta * torch.mean(
+ (hidden_state_quant - hidden_state.detach()) ** 2
+ )
+
+ # preserve gradients
+ hidden_state_quant = hidden_state + (hidden_state_quant - hidden_state).detach()
+
+ # reshape back to match original input shape
+ hidden_state_quant = hidden_state_quant.permute(0, 3, 1, 2).contiguous()
+
+ return hidden_state_quant, loss, min_encoding_indices
+
+
+class GlmImageVQVAEModelOutput(ChameleonVQVAEModelOutput):
+ pass
+
+
+class GlmImageVQVAE(ChameleonVQVAE):
+ _no_split_modules = [
+ "GlmImageVQVAEVectorQuantizer",
+ ]
+ _can_record_outputs = {}
+
+ def __init__(self, config: GlmImageVQVAEConfig):
+ super().__init__(config)
+ del self.encoder
+
+ def encode(self, hidden_states):
+ conv_hidden_states = self.quant_conv(hidden_states)
+ quantized_last_hidden_state, emb_loss, indices = self.quantize(conv_hidden_states)
+ return GlmImageVQVAEModelOutput(
+ last_hidden_state=hidden_states,
+ quantized_last_hidden_state=quantized_last_hidden_state,
+ image_tokens=indices,
+ embedding_loss=emb_loss,
+ )
+
+
+class GlmImageVisionModel(Glm4vVisionModel):
+ config: GlmImageVisionConfig
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+
+ def __init__(self, config: GlmImageVisionConfig):
+ super().__init__(config)
+
+ head_dim = config.hidden_size // config.num_heads
+ self.head_dim = head_dim
+
+ del self.merger
+ del self.rotary_pos_emb
+ del self.post_conv_layernorm
+ del self.downsample
+ del self.post_layernorm
+
+ def rot_pos_emb(self, grid_thw):
+ pos_ids = []
+ for t, h, w in grid_thw:
+ hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
+ hpos_ids = hpos_ids.reshape(
+ h // self.spatial_merge_size,
+ self.spatial_merge_size,
+ w // self.spatial_merge_size,
+ self.spatial_merge_size,
+ )
+ hpos_ids = hpos_ids.permute(0, 2, 1, 3)
+ hpos_ids = hpos_ids.flatten()
+
+ wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
+ wpos_ids = wpos_ids.reshape(
+ h // self.spatial_merge_size,
+ self.spatial_merge_size,
+ w // self.spatial_merge_size,
+ self.spatial_merge_size,
+ )
+ wpos_ids = wpos_ids.permute(0, 2, 1, 3)
+ wpos_ids = wpos_ids.flatten()
+ pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
+ pos_ids = torch.cat(pos_ids, dim=0)
+ return pos_ids
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self, pixel_values: torch.Tensor, grid_thw: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.Tensor` of shape `(total_patches, num_channels * patch_size * patch_size)`):
+ Packed pixel values.
+ grid_thw (`torch.Tensor` of shape `(num_images, 3)`):
+ The temporal, height and width of feature shape of each image.
+
+ Returns:
+ `torch.Tensor` of shape `(total_patches, hidden_size)`: Hidden states.
+ """
+
+ hidden_states = self.patch_embed(pixel_values)
+ image_type_ids = self.rot_pos_emb(grid_thw)
+
+ cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+ dim=0,
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+ )
+ cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+ seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()
+ hidden_states = self.embeddings(
+ hidden_states,
+ seqlens,
+ grid_thw,
+ image_type_ids[:, 0].to(hidden_states.device),
+ image_type_ids[:, 1].to(hidden_states.device),
+ )
+
+ # Transformer blocks (no position_embeddings needed, already added above)
+ for blk in self.blocks:
+ hidden_states = blk(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ )
+
+ return BaseModelOutputWithPooling(last_hidden_state=hidden_states)
+
+
+class GlmImageTextModel(Glm4vTextModel):
+ pass
+
+
+class GlmImageModel(Glm4vModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.visual = GlmImageVisionModel._from_config(config.vision_config)
+ self.language_model = GlmImageTextModel._from_config(config.text_config)
+ self.vqmodel = GlmImageVQVAE._from_config(config.vq_config)
+
+ self.rope_deltas = None # cache rope_deltas here
+
+ # Per-sample caches for batch processing
+ self._cached_decode_position_ids = None # shape: [batch_size, 3, max_decode_len]
+ self._prefill_len = None # prefill sequence length (same for all samples in batch)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_rope_index(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ images_per_sample: torch.LongTensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Calculate the 3D rope index for image generation task with full batch support.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary.
+ image_grid_thw (`torch.LongTensor` of shape `(total_images_in_batch, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image.
+ Images are packed across all samples in the batch.
+ images_per_sample (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Number of images (including target grids) for each sample in the batch.
+ Used to split image_grid_thw by sample.
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on padding token indices.
+
+ Returns:
+ position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`):
+ Position IDs for temporal, height, and width dimensions.
+ mrope_position_deltas (`torch.Tensor` of shape `(batch_size, 1)`):
+ Position deltas for multi-modal rotary position embedding.
+ """
+ batch_size, seq_len = input_ids.shape
+ device = input_ids.device
+ dtype = input_ids.dtype
+
+ image_start_token_id = self.config.image_start_token_id
+ image_end_token_id = self.config.image_end_token_id
+
+ position_ids = torch.ones(3, batch_size, seq_len, dtype=dtype, device=device)
+ text_positions = torch.arange(seq_len, device=device)[None, :].repeat(3, 1)
+
+ # Split image_grid_thw by sample if images_per_sample is provided
+ if image_grid_thw is not None and images_per_sample is not None:
+ grids_per_sample = torch.split(image_grid_thw, images_per_sample.tolist())
+ elif image_grid_thw is not None:
+ # Fallback: assume all grids belong to first sample (batch_size=1)
+ grids_per_sample = [image_grid_thw] * batch_size
+ else:
+ grids_per_sample = [None] * batch_size
+
+ # Per-sample caches for decode stage
+ all_decode_position_ids = []
+
+ for batch_idx in range(batch_size):
+ curr_input_ids = input_ids[batch_idx]
+ curr_grids = grids_per_sample[batch_idx]
+
+ if attention_mask is not None and attention_mask.shape[1] == seq_len:
+ valid_mask = attention_mask[batch_idx] == 1
+ curr_input_ids_valid = curr_input_ids[valid_mask]
+ else:
+ # attention_mask may have different length during assisted decoding
+ curr_input_ids_valid = curr_input_ids
+ valid_mask = None
+
+ # Find image boundaries in this sample
+ image_end_positions = torch.where(curr_input_ids_valid == image_end_token_id)[0]
+ image_start_positions = torch.where(curr_input_ids_valid == image_start_token_id)[0] + 1
+ num_complete_images = len(image_end_positions)
+
+ current_pos = 0
+ prev_image_end = 0
+ curr_position_ids = []
+
+ # Process complete images (source images in image-to-image task)
+ for img_idx, (start, end) in enumerate(zip(image_start_positions, image_end_positions)):
+ if curr_grids is None or img_idx >= len(curr_grids):
+ break
+
+ # Text tokens before this image
+ llm_pos_length = start - prev_image_end
+ llm_position_ids = text_positions[:, current_pos : current_pos + llm_pos_length].to(device=device)
+ current_pos += llm_position_ids.shape[-1]
+
+ # Image tokens with 2D spatial encoding
+ # For an image with height H and width W:
+ # - position_width cycles [0, 1, ..., W-1] for each row, repeated H times
+ # - position_height stays constant per row, [0]*W, [1]*W, ..., [H-1]*W
+ vision_position_ids = self.get_vision_position_ids(
+ start_position=current_pos, grid_thw=curr_grids[img_idx], device=device
+ )
+ current_pos += max(curr_grids[img_idx][1], curr_grids[img_idx][2])
+
+ prev_image_end = end
+ curr_position_ids.append(torch.cat([llm_position_ids, vision_position_ids], dim=-1))
+
+ # Remaining text tokens (including the final image_start token for generation)
+ end_position = len(curr_input_ids_valid) - prev_image_end
+ llm_position_ids = text_positions[:, current_pos : current_pos + end_position].to(device=device)
+ current_pos += llm_position_ids.shape[-1]
+ curr_position_ids.append(llm_position_ids)
+
+ # Concatenate all position ids for this sample
+ curr_position_ids = torch.cat(curr_position_ids, dim=-1)
+
+ # Store in the main position_ids tensor
+ if valid_mask is not None:
+ position_ids[:, batch_idx, valid_mask] = curr_position_ids
+ else:
+ position_ids[:, batch_idx, :] = curr_position_ids
+
+ # Build decode position ids for this sample
+ if curr_grids is not None and len(curr_grids) > 0:
+ num_decode_grids = len(curr_grids) - num_complete_images
+ num_decode_grids = max(num_decode_grids, 0)
+ decode_pos = current_pos
+
+ decode_temporal_list = []
+ decode_height_list = []
+ decode_width_list = []
+
+ curr_grids_list = curr_grids.tolist()
+ for i in range(1, num_decode_grids + 1):
+ grid_idx = -i
+ h = curr_grids_list[grid_idx][1]
+ w = curr_grids_list[grid_idx][2]
+ total_tokens = h * w
+
+ h_indices = torch.arange(h, device=device).unsqueeze(1).expand(h, w).flatten()
+ w_indices = torch.arange(w, device=device).unsqueeze(0).expand(h, w).flatten()
+
+ decode_temporal_list.append(
+ torch.full((total_tokens,), decode_pos, device=device, dtype=torch.long)
+ )
+ decode_height_list.append(decode_pos + h_indices)
+ decode_width_list.append(decode_pos + w_indices)
+ decode_pos = decode_pos + max(h, w)
+
+ # End marker
+ decode_temporal_list.append(torch.tensor([decode_pos], device=device, dtype=torch.long))
+ decode_height_list.append(torch.tensor([decode_pos], device=device, dtype=torch.long))
+ decode_width_list.append(torch.tensor([decode_pos], device=device, dtype=torch.long))
+
+ sample_decode_pos_ids = torch.stack(
+ [
+ torch.cat(decode_temporal_list, dim=0),
+ torch.cat(decode_height_list, dim=0),
+ torch.cat(decode_width_list, dim=0),
+ ],
+ dim=0,
+ )
+ all_decode_position_ids.append(sample_decode_pos_ids)
+
+ # Store prefill length (same for all samples since input_ids is padded to same length)
+ self._prefill_len = seq_len
+
+ # Pad decode position ids to same length and stack
+ if all_decode_position_ids:
+ max_decode_len = max(x.shape[1] for x in all_decode_position_ids)
+ padded_decode_pos_ids = [
+ F.pad(pos_ids, (0, max_decode_len - pos_ids.shape[1]), mode="replicate")
+ for pos_ids in all_decode_position_ids
+ ]
+ self._cached_decode_position_ids = torch.stack(padded_decode_pos_ids, dim=0) # [batch, 3, max_decode_len]
+ else:
+ self._cached_decode_position_ids = None
+
+ mrope_position_deltas = torch.zeros([batch_size, 1], dtype=dtype, device=device)
+
+ return position_ids, mrope_position_deltas
+
+ def get_image_tokens(
+ self,
+ hidden_states: torch.FloatTensor,
+ image_grid_thw: torch.LongTensor,
+ ) -> torch.LongTensor:
+ """
+ Tokenizes image features into discrete tokens with VQVAE module.
+
+ Args:
+ hidden_states (`torch.FloatTensor` of shape `(total_patches, hidden_size)`):
+ The packed image features from vision encoder.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`):
+ The temporal, height and width of feature shape of each image.
+
+ Returns:
+ image_tokens (`torch.LongTensor` of shape `(total_patches,)`):
+ Discrete token indices from the VQVAE codebook.
+ """
+ hidden_size = hidden_states.shape[-1]
+ split_sizes = (image_grid_thw.prod(dim=-1)).tolist()
+ hidden_states_list = torch.split(hidden_states, split_sizes, dim=0)
+
+ all_image_toks = []
+ for i, hs in enumerate(hidden_states_list):
+ grid_t, grid_h, grid_w = image_grid_thw[i].tolist()
+ hs = hs.view(grid_t, grid_h, grid_w, hidden_size)
+ hs = hs.permute(0, 3, 1, 2).contiguous()
+ vqmodel_outputs: GlmImageVQVAEModelOutput = self.vqmodel.encode(hs)
+ all_image_toks.append(vqmodel_outputs.image_tokens)
+ return torch.cat(all_image_toks, dim=0)
+
+ def get_video_features(self):
+ raise AttributeError("Not needed for GlmImage")
+
+ @can_return_tuple
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ pixel_values = pixel_values.type(self.visual.dtype)
+ vision_outputs = self.visual(pixel_values, grid_thw=image_grid_thw, return_dict=True, **kwargs)
+ split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
+ image_embeds = torch.split(vision_outputs.last_hidden_state, split_sizes)
+ vision_outputs.pooler_output = image_embeds
+
+ return vision_outputs
+
+ def get_placeholder_mask(
+ self,
+ input_ids: torch.LongTensor,
+ image_ids: torch.LongTensor,
+ ):
+ """
+ Replace image placeholder tokens in input_ids with actual image token ids from VQVAE.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, seq_len)`):
+ Input token ids with image placeholders.
+ image_ids (`torch.LongTensor` of shape `(num_images, num_tokens_per_image)` or flattened):
+ Discrete token indices from the VQVAE codebook.
+
+ Returns:
+ special_image_mask (`torch.LongTensor` of shape `(batch_size, seq_len)`):
+ Mask indicating positions in input ids that will be replaced by actual image tokens.
+ """
+
+ special_image_mask = input_ids == self.config.image_token_id
+ n_placeholder_tokens = special_image_mask.sum()
+ n_image_tokens = image_ids.shape[0]
+
+ if n_placeholder_tokens != n_image_tokens:
+ raise ValueError(
+ f"Number of image placeholder tokens ({n_placeholder_tokens.item()}) does not match "
+ f"number of image tokens from VQVAE ({n_image_tokens})"
+ )
+
+ return special_image_mask
+
+ def compute_3d_position_ids(
+ self,
+ input_ids: torch.Tensor | None,
+ image_grid_thw: torch.Tensor | None,
+ images_per_sample: torch.Tensor | None,
+ inputs_embeds: torch.Tensor | None,
+ attention_mask: torch.Tensor | None,
+ past_key_values: torch.Tensor | None,
+ ) -> torch.Tensor | None:
+ past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length()
+ can_compute_mrope = input_ids is not None and image_grid_thw is not None
+
+ if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0):
+ position_ids, rope_deltas = self.get_rope_index(
+ input_ids,
+ image_grid_thw=image_grid_thw,
+ attention_mask=attention_mask,
+ images_per_sample=images_per_sample,
+ )
+ self.rope_deltas = rope_deltas
+ # Use pre-calculated rope-deltas to infer correct 3D position ids during incremental
+ # generation (past_key_values_length > 0) or when only inputs_embeds is provided (no input_ids
+ # to recompute from). Skip when input_ids is provided without past_key_values to avoid shape
+ # mismatches from stale rope_deltas (e.g., training forward pass after generation).
+ elif self.rope_deltas is not None and (past_key_values_length > 0 or input_ids is None):
+ batch_size, seq_length, _ = inputs_embeds.shape
+ if self._cached_decode_position_ids is not None:
+ step = past_key_values_length - self._prefill_len
+ position_ids = self._cached_decode_position_ids[:, :, step : step + seq_length].permute(1, 0, 2)
+ else:
+ position_ids = (
+ torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_key_values_length
+ )
+ position_ids = position_ids.view(1, 1, -1).repeat(3, batch_size, 1)
+ else:
+ # Can't build correct 3D positions. Let the model infer it
+ position_ids = None
+ return position_ids
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ pixel_values: torch.Tensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ images_per_sample: torch.LongTensor | None = None,
+ rope_deltas: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | GlmImageModelOutputWithPast:
+ r"""
+ image_grid_thw (`torch.LongTensor` of shape `(total_images_in_batch, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ Images are packed across all samples in the batch.
+ images_per_sample (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Number of images (including target grids) for each sample in the batch.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ batch_size = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
+
+ if pixel_values is not None:
+ # Process source images (image-to-image mode)
+ # Source images are identified by counting image_end_token_id in input_ids
+ # Note: We must exclude padding tokens since pad_token_id == image_end_token_id
+ if images_per_sample is not None:
+ grids_per_sample = torch.split(image_grid_thw, images_per_sample.tolist())
+ # Create mask for non-padding tokens (attention_mask=1 means non-padding)
+ # Handle 4D attention mask (from static cache) by extracting diagonal
+ if attention_mask is not None and attention_mask.ndim == 4:
+ non_pad_mask = torch.diagonal(attention_mask[:, 0], dim1=1, dim2=2)
+ if non_pad_mask.dtype.is_floating_point:
+ non_pad_mask = non_pad_mask / torch.finfo(non_pad_mask.dtype).min
+ non_pad_mask = (1.0 - non_pad_mask).int()
+ # Only keep columns matching input_ids length
+ non_pad_mask = non_pad_mask[:, -input_ids.shape[1] :]
+ else:
+ non_pad_mask = attention_mask if attention_mask is not None else torch.ones_like(input_ids)
+
+ source_grids_list = []
+ is_image_end = input_ids == self.config.image_end_token_id
+ is_non_pad = non_pad_mask == 1
+ num_source_per_sample = (is_image_end & is_non_pad).sum(dim=1).tolist()
+ for sample_idx in range(batch_size):
+ num_source = num_source_per_sample[sample_idx]
+ if num_source > 0:
+ source_grids_list.append(grids_per_sample[sample_idx][:num_source])
+ if len(source_grids_list) == 0:
+ raise ValueError(
+ "pixel_values provided but no source images found in input_ids. "
+ "Ensure input_ids contains image_end_token_id for each source image."
+ )
+ source_grids = torch.cat(source_grids_list, dim=0)
+ else:
+ # Fallback for batch_size=1: all but last grid are source images
+ source_grids = image_grid_thw[:-1]
+
+ image_features = self.get_image_features(pixel_values, source_grids, return_dict=True)
+ image_embeds = torch.cat(image_features.pooler_output, dim=0)
+ image_ids = self.get_image_tokens(image_embeds, source_grids)
+ image_ids = image_ids.view(-1).to(input_ids.device)
+ special_image_mask = self.get_placeholder_mask(input_ids, image_ids)
+ input_ids = input_ids.masked_scatter(special_image_mask, image_ids)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if position_ids is None:
+ position_ids = self.compute_3d_position_ids(
+ input_ids=input_ids,
+ image_grid_thw=image_grid_thw,
+ images_per_sample=images_per_sample,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ )
+
+ outputs = self.language_model(
+ input_ids=None,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ return GlmImageModelOutputWithPast(
+ last_hidden_state=outputs.last_hidden_state,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ rope_deltas=self.rope_deltas,
+ )
+
+
+class GlmImageCausalLMOutputWithPast(Glm4vCausalLMOutputWithPast):
+ pass
+
+
+class GlmImageForConditionalGeneration(GlmImagePreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {}
+ # Reference: fix gemma3 grad acc #37208
+ accepts_loss_kwargs = False
+ base_model_prefix = "model"
+ config: GlmImageConfig
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = GlmImageModel(config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vision_vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ return self.model.get_image_features(pixel_values, image_grid_thw, **kwargs)
+
+ def get_image_tokens(self, hidden_states: torch.FloatTensor, image_grid_thw: torch.LongTensor | None = None):
+ return self.model.get_image_tokens(hidden_states, image_grid_thw)
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ pixel_values: torch.Tensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ images_per_sample: torch.LongTensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | GlmImageCausalLMOutputWithPast:
+ r"""
+ 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]`.
+ image_grid_thw (`torch.LongTensor` of shape `(total_images_in_batch, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ Images are packed across all samples in the batch.
+ images_per_sample (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Number of images (including target grids) for each sample in the batch.
+
+ Example:
+
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoProcessor, GlmImageForConditionalGeneration
+
+ >>> model = GlmImageForConditionalGeneration.from_pretrained("zai-org/GLM-Image")
+ >>> processor = AutoProcessor.from_pretrained("zai-org/GLM-Image")
+
+ >>> messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "image"},
+ {"type": "text", "text": "Add a truck of this photo.28 40"},
+ ],
+ },
+ ]
+ >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+ >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos])
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."
+ ```"""
+ outputs = self.model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ image_grid_thw=image_grid_thw,
+ images_per_sample=images_per_sample,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)
+
+ return GlmImageCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ rope_deltas=outputs.rope_deltas,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ image_grid_thw=None,
+ images_per_sample=None,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ position_ids=position_ids,
+ pixel_values=pixel_values,
+ image_grid_thw=image_grid_thw,
+ is_first_iteration=is_first_iteration,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ model_inputs["position_ids"] = None
+ model_inputs["images_per_sample"] = images_per_sample
+
+ if not is_first_iteration and use_cache:
+ model_inputs["pixel_values"] = None
+
+ return model_inputs
+
+ def _get_image_nums(
+ self,
+ input_ids: torch.LongTensor | None,
+ ) -> torch.Tensor:
+ """
+ Get the number of images for each sample.
+ For GLM-Image, only input_ids allow us to get the number of images.
+
+ Returns:
+ image_counts (`torch.LongTensor` of shape `(batch_size,)`)
+ """
+ is_image = input_ids == self.config.image_start_token_id
+
+ return is_image.sum(dim=1)
+
+ def _expand_inputs_for_generation(
+ self,
+ expand_size: int = 1,
+ is_encoder_decoder: bool = False,
+ input_ids: torch.LongTensor | None = None,
+ **model_kwargs,
+ ) -> tuple[torch.LongTensor, dict[str, Any]]:
+ # Overwritten -- Support for expanding tensors without a batch size dimension
+ # e.g., pixel_values, image_grid_thw
+ # pixel_values.shape[0] is sum(seqlen_images for samples)
+ # image_grid_thw.shape[0] is sum(num_images for samples)
+
+ if expand_size == 1:
+ return input_ids, model_kwargs
+
+ visual_keys = ["pixel_values", "image_grid_thw", "images_per_sample"]
+
+ def _expand_dict_for_generation_visual(dict_to_expand):
+ image_grid_thw = model_kwargs.get("image_grid_thw", None)
+ if image_grid_thw is None:
+ return dict_to_expand
+
+ images_per_sample = model_kwargs.get("images_per_sample", None)
+
+ # Use images_per_sample if available
+ if images_per_sample is not None:
+ image_nums = images_per_sample.tolist()
+ elif input_ids is not None:
+ # Try to infer from image_grid_thw / batch_size
+ batch_size = input_ids.shape[0]
+ total_grids = image_grid_thw.shape[0]
+ if total_grids % batch_size == 0:
+ grids_per_sample = total_grids // batch_size
+ image_nums = [grids_per_sample] * batch_size
+ else:
+ # Cannot evenly distribute grids - fall back to simple repeat_interleave
+ # This handles test cases where image_grid_thw has (batch_size + 1) rows
+ dict_to_expand["image_grid_thw"] = image_grid_thw.repeat_interleave(expand_size, dim=0)
+ if dict_to_expand.get("pixel_values") is not None:
+ dict_to_expand["pixel_values"] = dict_to_expand["pixel_values"].repeat_interleave(
+ expand_size, dim=0
+ )
+ return dict_to_expand
+ else:
+ image_nums = self._get_image_nums(input_ids).tolist()
+
+ # Get source image counts per sample from image_end_token_id count
+ source_image_nums = (input_ids == self.config.image_end_token_id).sum(dim=1).tolist()
+
+ def _repeat_interleave_samples(x, lengths, repeat_times):
+ samples = torch.split(x, lengths)
+ repeat_args = [repeat_times] + [1] * (x.dim() - 1)
+ result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0)
+ return result
+
+ for key in dict_to_expand:
+ if key == "pixel_values":
+ # Split images into samples based on source image counts
+ if sum(source_image_nums) > 0:
+ # Split grids by sample to compute pixel counts
+ grids_per_sample = torch.split(image_grid_thw, image_nums)
+ all_pixel_counts = image_grid_thw.prod(dim=1)
+ pixel_counts_per_sample = torch.split(all_pixel_counts, image_nums)
+ # Build source mask and compute per-sample source pixel counts in one sync
+ source_pixel_counts = torch.zeros(len(grids_per_sample), device=image_grid_thw.device)
+ for batch_idx in range(len(grids_per_sample)):
+ num_source = source_image_nums[batch_idx]
+ if num_source > 0:
+ source_pixel_counts[batch_idx] = pixel_counts_per_sample[batch_idx][:num_source].sum()
+ lengths = source_pixel_counts.to(torch.int64).tolist()
+
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "image_grid_thw":
+ # Expand all grids (source + target) per sample
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=image_nums, repeat_times=expand_size
+ )
+ elif key == "images_per_sample":
+ # Simply repeat the counts
+ if dict_to_expand.get(key) is not None:
+ dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
+ return dict_to_expand
+
+ def _expand_dict_for_generation(dict_to_expand):
+ for key in dict_to_expand:
+ if (
+ dict_to_expand[key] is not None
+ and isinstance(dict_to_expand[key], torch.Tensor)
+ and key not in visual_keys
+ ):
+ dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
+ return dict_to_expand
+
+ model_kwargs = _expand_dict_for_generation_visual(model_kwargs)
+
+ if input_ids is not None:
+ input_ids = input_ids.repeat_interleave(expand_size, dim=0)
+
+ model_kwargs = _expand_dict_for_generation(model_kwargs)
+
+ if is_encoder_decoder:
+ if model_kwargs.get("encoder_outputs") is None:
+ raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.")
+ model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])
+
+ return input_ids, model_kwargs
+
+
+def smart_resize(
+ height: int,
+ width: int,
+ factor: int = 16,
+ min_pixels: int = 512 * 512,
+ max_pixels: int = 2048 * 2048,
+) -> tuple[int, int]:
+ if height < factor or width < factor:
+ raise ValueError(f"height:{height} or width:{width} must be larger than factor:{factor}")
+ elif max(height, width) / min(height, width) > 4:
+ raise ValueError(
+ f"absolute aspect ratio must be smaller than 4, got {max(height, width) / min(height, width)}"
+ )
+
+ shortest_edge = int(round(math.sqrt(min_pixels)))
+ longest_edge = int(round(math.sqrt(max_pixels)))
+ min_side = min(height, width)
+ max_side = max(height, width)
+
+ scale = 1.0
+
+ if min_side < shortest_edge:
+ scale = shortest_edge / min_side
+
+ if max_side * scale > longest_edge:
+ scale = longest_edge / max_side
+
+ height = height // 2
+ width = width // 2
+
+ h_bar = max(factor, int(round(height * scale / factor)) * factor)
+ w_bar = max(factor, int(round(width * scale / factor)) * factor)
+
+ if max(h_bar, w_bar) > longest_edge:
+ beta = max(h_bar, w_bar) / longest_edge
+ h_bar = max(factor, int(math.floor((h_bar / beta) / factor)) * factor)
+ w_bar = max(factor, int(math.floor((w_bar / beta) / factor)) * factor)
+
+ return h_bar, w_bar
+
+
+class GlmImageImageProcessor(Qwen2VLImageProcessor):
+ model_input_names = ["pixel_values", "image_grid_thw", "images_per_sample"]
+
+
+class GlmImageImageProcessorPil(Qwen2VLImageProcessorPil):
+ model_input_names = ["pixel_values", "image_grid_thw", "images_per_sample"]
+
+
+class GlmImageImagesKwargs(ImagesKwargs, total=False):
+ """
+ target_h (`int`):
+ Height of the target image to be generated.
+ target_w (`int`):
+ Width of the target image to be generated.
+ """
+
+ target_h: int
+ target_w: int
+
+
+class GlmImageProcessorKwargs(Qwen2VLProcessorKwargs):
+ images_kwargs: GlmImageImagesKwargs
+
+ _defaults = {
+ "text_kwargs": {
+ "padding": False,
+ "return_mm_token_type_ids": False,
+ },
+ "images_kwargs": {
+ "target_h": 1152,
+ "target_w": 768,
+ },
+ }
+
+
+@requires(backends=("torch",))
+class GlmImageProcessor(ProcessorMixin):
+ r"""
+ Constructs a GLM-Image processor which wraps a GLM-Image image processor and a GLM-Image tokenizer into a single processor.
+ [`~GlmImageProcessor.__call__`] and [`~GlmImageProcessor.decode`] for more information.
+ Args:
+ image_processor ([`GlmImageProcessor`], *optional*):
+ The image processor is a required input.
+ tokenizer ([`PreTrainedTokenizerFast`], *optional*):
+ The tokenizer is a required input.
+ chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
+ in a chat into a tokenizable string.
+ """
+
+ model_input_names = ["input_ids", "attention_mask", "pixel_values", "image_grid_thw", "images_per_sample"]
+
+ def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs):
+ self.image_token = tokenizer.image_token
+ self.grid_bos_token = tokenizer.grid_bos_token
+ self.grid_eos_token = tokenizer.grid_eos_token
+ self.bos_token = tokenizer.bos_token
+ self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token)
+ super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+ def __call__(
+ self,
+ images: ImageInput | None = None,
+ text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None,
+ **kwargs: Unpack[GlmImageProcessorKwargs],
+ ) -> BatchFeature:
+ """
+ Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
+ and `kwargs` arguments to PreTrainedTokenizerFast's [`~PreTrainedTokenizerFast.__call__`] if `text` is not `None` to encode
+ the text.
+
+ Args:
+ images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
+ The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+ tensor. Both channels-first and channels-last formats are supported.
+ text (`str`, `List[str]`, `List[List[str]]`):
+ The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+ (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+ `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+ return_tensors (`str` or [`~utils.TensorType`], *optional*):
+ If set, will return tensors of a particular framework. Acceptable values are:
+ - `'pt'`: Return PyTorch `torch.Tensor` objects.
+ - `'np'`: Return NumPy `np.ndarray` objects.
+
+ Returns:
+ [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+ - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+ - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+ `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+ `None`).
+ - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+ - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`.
+ """
+ output_kwargs = self._merge_kwargs(
+ GlmImageProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+
+ target_h = output_kwargs["images_kwargs"].pop("target_h", None)
+ target_w = output_kwargs["images_kwargs"].pop("target_w", None)
+ is_text_to_image = images is None
+
+ if images is not None:
+ image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
+ image_grid_thw = image_inputs["image_grid_thw"]
+ else:
+ image_inputs = {}
+ image_grid_thw = None
+
+ # Handle text=None case (image-only processing)
+ if text is None:
+ if images is None:
+ raise ValueError("You must provide at least one of `text` or `images`.")
+ return image_inputs
+
+ if not isinstance(text, list):
+ text = [text]
+
+ batch_size = len(text)
+ text = text.copy() # below lines change text in-place
+
+ # Count images per sample by counting image tokens in each text
+ images_per_sample = []
+ for i in range(batch_size):
+ images_per_sample.append(text[i].count(self.image_token))
+
+ # Replace image tokens with the correct number of placeholder tokens
+ if not is_text_to_image:
+ index = 0
+ for i in range(batch_size):
+ while self.image_token in text[i]:
+ grid = image_grid_thw[index]
+ num_image_tokens = int(grid[1] * grid[2])
+ text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_image_tokens, 1)
+ index += 1
+ text[i] = text[i].replace("<|placeholder|>", self.image_token)
+
+ # Build prompt with target shape and combine grids in a single loop
+ # Format: [sample0_source_grids..., sample0_target_grids, sample1_source_grids..., sample1_target_grids, ...]
+ # Note: In i2i mode, batches are homogeneous (same number of source images per sample)
+ num_source_images = images_per_sample[0] if images_per_sample else 0
+
+ # Validate homogeneity for i2i mode
+ if not is_text_to_image and images_per_sample and len(set(images_per_sample)) != 1:
+ raise ValueError(
+ f"In image-to-image mode, all samples must have the same number of source images. "
+ f"Got different counts: {images_per_sample}"
+ )
+
+ all_grids = []
+ for i in range(batch_size):
+ text[i], token_h, token_w, prev_h, prev_w = self._build_prompt_with_target_shape(
+ text[i], height=target_h, width=target_w, is_text_to_image=is_text_to_image
+ )
+ # Add source grids for this sample (i2i mode only)
+ if not is_text_to_image and num_source_images > 0:
+ start_idx = i * num_source_images
+ all_grids.append(image_grid_thw[start_idx : start_idx + num_source_images])
+ # Add target grid for this sample
+ all_grids.append(
+ self._build_target_image_grid_thw(
+ token_h=token_h,
+ token_w=token_w,
+ prev_token_h=prev_h,
+ prev_token_w=prev_w,
+ is_text_to_image=is_text_to_image,
+ )
+ )
+ image_inputs["image_grid_thw"] = torch.cat(all_grids, dim=0)
+
+ # Store images_per_sample for later use (add target images count)
+ # Each sample will have: source_images + target_images (typically 2 for t2i, 1 for i2i)
+ num_target_grids = 2 if is_text_to_image else 1
+ image_inputs["images_per_sample"] = torch.tensor(
+ [num_source_images + num_target_grids] * batch_size, dtype=torch.long
+ )
+
+ return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
+ return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False)
+ text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+
+ self._check_special_mm_tokens(text, text_inputs, modalities=["image"])
+
+ if return_mm_token_type_ids:
+ text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"])
+ return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors)
+
+ def _build_prompt_with_target_shape(
+ self,
+ prompt: str,
+ height: int,
+ width: int,
+ is_text_to_image: bool,
+ ) -> tuple[str, int, int, int, int]:
+ factor = 32
+ height = (height // factor) * factor
+ width = (width // factor) * factor
+ token_h = height // factor
+ token_w = width // factor
+ ratio = token_h / token_w
+ prev_token_h = int(math.sqrt(ratio) * (factor // 2))
+ prev_token_w = int(math.sqrt(1 / ratio) * (factor // 2))
+
+ if is_text_to_image:
+ expanded_prompt = f"{prompt}{self.grid_bos_token}{token_h} {token_w}{self.grid_eos_token}{self.grid_bos_token}{prev_token_h} {prev_token_w}{self.grid_eos_token}{self.bos_token}"
+ else:
+ expanded_prompt = f"{prompt}{self.grid_bos_token}{token_h} {token_w}{self.grid_eos_token}{self.bos_token}"
+
+ return expanded_prompt, token_h, token_w, prev_token_h, prev_token_w
+
+ @staticmethod
+ def _build_target_image_grid_thw(
+ token_h: int,
+ token_w: int,
+ prev_token_h: int,
+ prev_token_w: int,
+ is_text_to_image: bool = True,
+ ):
+ if is_text_to_image:
+ # Text-to-image: 2 target grids (large + small preview)
+ return torch.tensor(
+ [
+ [1, token_h, token_w],
+ [1, prev_token_h, prev_token_w],
+ ],
+ )
+ else:
+ # Image-to-image: 1 target grid only
+ return torch.tensor(
+ [
+ [1, token_h, token_w],
+ ],
+ )
+
+
+__all__ = [
+ "GlmImageVQVAEConfig",
+ "GlmImageVisionConfig",
+ "GlmImageTextConfig",
+ "GlmImageConfig",
+ "GlmImagePreTrainedModel",
+ "GlmImageVQVAE",
+ "GlmImageVisionModel",
+ "GlmImageTextModel",
+ "GlmImageModel",
+ "GlmImageForConditionalGeneration",
+ "GlmImageImageProcessor",
+ "GlmImageImageProcessorPil",
+ "GlmImageProcessor",
+]
diff --git a/third_party/transformers/src/transformers/models/glm_image/processing_glm_image.py b/third_party/transformers/src/transformers/models/glm_image/processing_glm_image.py
new file mode 100644
index 0000000000000000000000000000000000000000..1b117f2ced4ea47e5b31293256ac2b7120b89cc0
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_image/processing_glm_image.py
@@ -0,0 +1,259 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/glm_image/modular_glm_image.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_glm_image.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2025 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+import math
+
+import torch
+
+from ...feature_extraction_utils import BatchFeature
+from ...image_utils import ImageInput
+from ...processing_utils import ImagesKwargs, ProcessingKwargs, ProcessorMixin, Unpack
+from ...tokenization_utils_base import PreTokenizedInput, TextInput
+from ...utils.import_utils import requires
+
+
+class GlmImageImagesKwargs(ImagesKwargs, total=False):
+ """
+ target_h (`int`):
+ Height of the target image to be generated.
+ target_w (`int`):
+ Width of the target image to be generated.
+ """
+
+ target_h: int
+ target_w: int
+
+
+class GlmImageProcessorKwargs(ProcessingKwargs, total=False):
+ _defaults = {
+ "text_kwargs": {
+ "padding": False,
+ "return_mm_token_type_ids": False,
+ },
+ "images_kwargs": {
+ "target_h": 1152,
+ "target_w": 768,
+ },
+ }
+ images_kwargs: GlmImageImagesKwargs
+
+
+@requires(backends=("torch",))
+class GlmImageProcessor(ProcessorMixin):
+ r"""
+ Constructs a GLM-Image processor which wraps a GLM-Image image processor and a GLM-Image tokenizer into a single processor.
+ [`~GlmImageProcessor.__call__`] and [`~GlmImageProcessor.decode`] for more information.
+ Args:
+ image_processor ([`GlmImageProcessor`], *optional*):
+ The image processor is a required input.
+ tokenizer ([`PreTrainedTokenizerFast`], *optional*):
+ The tokenizer is a required input.
+ chat_template (`str`, *optional*): A Jinja template which will be used to convert lists of messages
+ in a chat into a tokenizable string.
+ """
+
+ model_input_names = ["input_ids", "attention_mask", "pixel_values", "image_grid_thw", "images_per_sample"]
+
+ def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs):
+ self.image_token = tokenizer.image_token
+ self.grid_bos_token = tokenizer.grid_bos_token
+ self.grid_eos_token = tokenizer.grid_eos_token
+ self.bos_token = tokenizer.bos_token
+ self.image_token_id = tokenizer.convert_tokens_to_ids(self.image_token)
+ super().__init__(image_processor, tokenizer, chat_template=chat_template)
+
+ def __call__(
+ self,
+ images: ImageInput | None = None,
+ text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None,
+ **kwargs: Unpack[GlmImageProcessorKwargs],
+ ) -> BatchFeature:
+ """
+ Main method to prepare for the model one or several sequences(s) and image(s). This method forwards the `text`
+ and `kwargs` arguments to PreTrainedTokenizerFast's [`~PreTrainedTokenizerFast.__call__`] if `text` is not `None` to encode
+ the text.
+
+ Args:
+ images (`PIL.Image.Image`, `np.ndarray`, `torch.Tensor`, `List[PIL.Image.Image]`, `List[np.ndarray]`, `List[torch.Tensor]`):
+ The image or batch of images to be prepared. Each image can be a PIL image, NumPy array or PyTorch
+ tensor. Both channels-first and channels-last formats are supported.
+ text (`str`, `List[str]`, `List[List[str]]`):
+ The sequence or batch of sequences to be encoded. Each sequence can be a string or a list of strings
+ (pretokenized string). If the sequences are provided as list of strings (pretokenized), you must set
+ `is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
+ return_tensors (`str` or [`~utils.TensorType`], *optional*):
+ If set, will return tensors of a particular framework. Acceptable values are:
+ - `'pt'`: Return PyTorch `torch.Tensor` objects.
+ - `'np'`: Return NumPy `np.ndarray` objects.
+
+ Returns:
+ [`BatchFeature`]: A [`BatchFeature`] with the following fields:
+
+ - **input_ids** -- List of token ids to be fed to a model. Returned when `text` is not `None`.
+ - **attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+ `return_attention_mask=True` or if *"attention_mask"* is in `self.model_input_names` and if `text` is not
+ `None`).
+ - **pixel_values** -- Pixel values to be fed to a model. Returned when `images` is not `None`.
+ - **image_grid_thw** -- List of image 3D grid in LLM. Returned when `images` is not `None`.
+ """
+ output_kwargs = self._merge_kwargs(
+ GlmImageProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+
+ target_h = output_kwargs["images_kwargs"].pop("target_h", None)
+ target_w = output_kwargs["images_kwargs"].pop("target_w", None)
+ is_text_to_image = images is None
+
+ if images is not None:
+ image_inputs = self.image_processor(images=images, **output_kwargs["images_kwargs"])
+ image_grid_thw = image_inputs["image_grid_thw"]
+ else:
+ image_inputs = {}
+ image_grid_thw = None
+
+ # Handle text=None case (image-only processing)
+ if text is None:
+ if images is None:
+ raise ValueError("You must provide at least one of `text` or `images`.")
+ return image_inputs
+
+ if not isinstance(text, list):
+ text = [text]
+
+ batch_size = len(text)
+ text = text.copy() # below lines change text in-place
+
+ # Count images per sample by counting image tokens in each text
+ images_per_sample = []
+ for i in range(batch_size):
+ images_per_sample.append(text[i].count(self.image_token))
+
+ # Replace image tokens with the correct number of placeholder tokens
+ if not is_text_to_image:
+ index = 0
+ for i in range(batch_size):
+ while self.image_token in text[i]:
+ grid = image_grid_thw[index]
+ num_image_tokens = int(grid[1] * grid[2])
+ text[i] = text[i].replace(self.image_token, "<|placeholder|>" * num_image_tokens, 1)
+ index += 1
+ text[i] = text[i].replace("<|placeholder|>", self.image_token)
+
+ # Build prompt with target shape and combine grids in a single loop
+ # Format: [sample0_source_grids..., sample0_target_grids, sample1_source_grids..., sample1_target_grids, ...]
+ # Note: In i2i mode, batches are homogeneous (same number of source images per sample)
+ num_source_images = images_per_sample[0] if images_per_sample else 0
+
+ # Validate homogeneity for i2i mode
+ if not is_text_to_image and images_per_sample and len(set(images_per_sample)) != 1:
+ raise ValueError(
+ f"In image-to-image mode, all samples must have the same number of source images. "
+ f"Got different counts: {images_per_sample}"
+ )
+
+ all_grids = []
+ for i in range(batch_size):
+ text[i], token_h, token_w, prev_h, prev_w = self._build_prompt_with_target_shape(
+ text[i], height=target_h, width=target_w, is_text_to_image=is_text_to_image
+ )
+ # Add source grids for this sample (i2i mode only)
+ if not is_text_to_image and num_source_images > 0:
+ start_idx = i * num_source_images
+ all_grids.append(image_grid_thw[start_idx : start_idx + num_source_images])
+ # Add target grid for this sample
+ all_grids.append(
+ self._build_target_image_grid_thw(
+ token_h=token_h,
+ token_w=token_w,
+ prev_token_h=prev_h,
+ prev_token_w=prev_w,
+ is_text_to_image=is_text_to_image,
+ )
+ )
+ image_inputs["image_grid_thw"] = torch.cat(all_grids, dim=0)
+
+ # Store images_per_sample for later use (add target images count)
+ # Each sample will have: source_images + target_images (typically 2 for t2i, 1 for i2i)
+ num_target_grids = 2 if is_text_to_image else 1
+ image_inputs["images_per_sample"] = torch.tensor(
+ [num_source_images + num_target_grids] * batch_size, dtype=torch.long
+ )
+
+ return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
+ return_mm_token_type_ids = output_kwargs["text_kwargs"].pop("return_mm_token_type_ids", False)
+ text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"])
+
+ self._check_special_mm_tokens(text, text_inputs, modalities=["image"])
+
+ if return_mm_token_type_ids:
+ text_inputs["mm_token_type_ids"] = self.create_mm_token_type_ids(text_inputs["input_ids"])
+ return BatchFeature(data={**text_inputs, **image_inputs}, tensor_type=return_tensors)
+
+ def _build_prompt_with_target_shape(
+ self,
+ prompt: str,
+ height: int,
+ width: int,
+ is_text_to_image: bool,
+ ) -> tuple[str, int, int, int, int]:
+ factor = 32
+ height = (height // factor) * factor
+ width = (width // factor) * factor
+ token_h = height // factor
+ token_w = width // factor
+ ratio = token_h / token_w
+ prev_token_h = int(math.sqrt(ratio) * (factor // 2))
+ prev_token_w = int(math.sqrt(1 / ratio) * (factor // 2))
+
+ if is_text_to_image:
+ expanded_prompt = f"{prompt}{self.grid_bos_token}{token_h} {token_w}{self.grid_eos_token}{self.grid_bos_token}{prev_token_h} {prev_token_w}{self.grid_eos_token}{self.bos_token}"
+ else:
+ expanded_prompt = f"{prompt}{self.grid_bos_token}{token_h} {token_w}{self.grid_eos_token}{self.bos_token}"
+
+ return expanded_prompt, token_h, token_w, prev_token_h, prev_token_w
+
+ @staticmethod
+ def _build_target_image_grid_thw(
+ token_h: int,
+ token_w: int,
+ prev_token_h: int,
+ prev_token_w: int,
+ is_text_to_image: bool = True,
+ ):
+ if is_text_to_image:
+ # Text-to-image: 2 target grids (large + small preview)
+ return torch.tensor(
+ [
+ [1, token_h, token_w],
+ [1, prev_token_h, prev_token_w],
+ ],
+ )
+ else:
+ # Image-to-image: 1 target grid only
+ return torch.tensor(
+ [
+ [1, token_h, token_w],
+ ],
+ )
+
+
+__all__ = ["GlmImageProcessor"]
diff --git a/third_party/transformers/src/transformers/models/glm_moe_dsa/__init__.py b/third_party/transformers/src/transformers/models/glm_moe_dsa/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d874c6c2183277d691d3dc134c9d340bffbaaa9c
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_moe_dsa/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2026 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_glm_moe_dsa import *
+ from .modeling_glm_moe_dsa import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py b/third_party/transformers/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f9f2a766cf4665cf6db63e0e984fab08e9060a0
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_moe_dsa/configuration_glm_moe_dsa.py
@@ -0,0 +1,132 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_glm_moe_dsa.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2026 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="zai-org/GLM-5")
+@strict
+class GlmMoeDsaConfig(PreTrainedConfig):
+ r"""
+ n_group (`int`, *optional*, defaults to 1):
+ Number of groups for routed experts.
+ mlp_layer_types (`list`, *optional*):
+ MLP type pattern for each layer (`"dense"` or `"sparse"`). Defaults to 3 dense + rest sparse.
+ index_topk (`int`, *optional*, defaults to 2048):
+ Number of top tokens selected by the indexer for sparse attention.
+ index_head_dim (`int`, *optional*, defaults to 128):
+ Head dimension for the indexer projections (DSA).
+ index_n_heads (`int | None`, *optional*, defaults to 32):
+ Number of heads for the indexer projections (DSA).
+
+ ```python
+ >>> from transformers import GlmMoeDsaConfig, GlmMoeDsaModel
+
+ >>> # Initializing a GLM-MoE-DSA configuration
+ >>> configuration = GlmMoeDsaConfig()
+
+ >>> # Initializing a model from the configuration
+ >>> model = GlmMoeDsaModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "glm_moe_dsa"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_b_proj": "colwise",
+ "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
+ "layers.*.self_attn.kv_b_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
+ "layers.*.mlp.experts.down_proj": "rowwise",
+ "layers.*.mlp.experts": "moe_tp_experts",
+ "layers.*.mlp.shared_experts.gate_proj": "colwise",
+ "layers.*.mlp.shared_experts.up_proj": "colwise",
+ "layers.*.mlp.shared_experts.down_proj": "rowwise",
+ "layers.*.mlp.gate_proj": "colwise",
+ "layers.*.mlp.up_proj": "colwise",
+ "layers.*.mlp.down_proj": "rowwise",
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+ attribute_map = {
+ "num_local_experts": "n_routed_experts",
+ "head_dim": "qk_rope_head_dim",
+ }
+
+ vocab_size: int = 154880
+
+ hidden_size: int = 6144
+ intermediate_size: int = 12288
+ moe_intermediate_size: int = 2048
+ num_hidden_layers: int = 78
+ num_attention_heads: int = 64
+ num_key_value_heads: int = 64
+ n_shared_experts: int = 1
+ n_routed_experts: int = 256
+ routed_scaling_factor: float = 2.5
+ kv_lora_rank: int = 512
+ q_lora_rank: int = 2048
+ qk_rope_head_dim: int = 64
+ v_head_dim: int = 256
+ qk_nope_head_dim: int = 192
+ n_group: int = 1
+ topk_group: int = 1
+ num_experts_per_tok: int = 8
+ norm_topk_prob: bool = True
+ hidden_act: str = "silu"
+ max_position_embeddings: int = 202752
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-5
+ use_cache: bool = True
+ pad_token_id: int | None = None
+ bos_token_id: int | None = 0
+ eos_token_id: int | list[int] | None = 1
+ tie_word_embeddings: bool = False
+ rope_parameters: RopeParameters | dict | None = None
+ mlp_layer_types: list[str] | None = None
+ attention_bias: bool = False
+ attention_dropout: float | int = 0.0
+ index_topk: int = 2048
+ index_head_dim: int = 128
+ index_n_heads: int = 32
+
+ def __post_init__(self, **kwargs):
+ self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+
+ # MLP layer types: first 3 dense, rest sparse
+ if self.mlp_layer_types is None:
+ self.mlp_layer_types = ["dense"] * min(3, self.num_hidden_layers) + ["sparse"] * (
+ self.num_hidden_layers - 3
+ )
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["GlmMoeDsaConfig"]
diff --git a/third_party/transformers/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py b/third_party/transformers/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py
new file mode 100644
index 0000000000000000000000000000000000000000..950deba0800e09a3960c50c4b45f7369d7d9f534
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_moe_dsa/modeling_glm_moe_dsa.py
@@ -0,0 +1,879 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_glm_moe_dsa.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2026 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_experts_implementation, use_kernel_forward_from_hub
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
+from ...utils.generic import is_flash_attention_requested, maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_glm_moe_dsa import GlmMoeDsaConfig
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class GlmMoeDsaRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ GlmMoeDsaRMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+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)
+
+
+def apply_rotary_pos_emb(
+ x: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+ unsqueeze_dim: int = 1,
+) -> torch.Tensor:
+ """
+ Applies Rotary Position Embedding to a single tensor.
+
+ This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved)`.
+ Instead of using complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding.
+
+ Args:
+ x (`torch.Tensor`): Input tensor of shape `[..., head_dim]`.
+ cos (`torch.Tensor`): Cosine part from RotaryEmbedding, shape `[batch, seq_len, head_dim]`.
+ sin (`torch.Tensor`): Sine part from RotaryEmbedding, shape `[batch, seq_len, head_dim]`.
+ unsqueeze_dim (`int`): Dimension along which to unsqueeze cos/sin for broadcasting.
+ Use `1` when x is `[B, H, S, D]` (BHSD) and `2` when x is `[B, S, H, D]` (BSHD).
+
+ Returns:
+ `torch.Tensor`: Tensor with rotary embeddings applied, same shape as input.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+
+ # Split-half (NeoX/Llama style): (x[:d/2], x[d/2:])
+ # This matches llama's apply_rotary_pos_emb logic.
+ x_rotated = (x * cos) + (rotate_half(x) * sin)
+ return x_rotated
+
+
+class GlmMoeDsaIndexer(nn.Module):
+ """
+ Dynamic Sparse Attention (DSA) indexer for selecting top-k tokens.
+
+ The Indexer has its own lightweight projections (wq_b, wk) separate from the
+ main MLA attention. It uses non-interleaved (NeoX/Llama) RoPE, unlike the main attention
+ which uses interleaved RoPE.
+
+ **Cache strategy**: The Indexer manages its own key cache (`_cached_keys`) separately
+ from the DynamicCache used by MLA attention, since DynamicCache is sized for exactly
+ `num_hidden_layers` attention layers. Keys are concatenated along the sequence dimension
+ during autoregressive decode.
+ """
+
+ def __init__(self, config: "GlmMoeDsaConfig", layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+
+ self.hidden_size: int = config.hidden_size
+ self.n_heads: int = config.index_n_heads
+ self.head_dim: int = config.index_head_dim
+ self.qk_rope_head_dim: int = config.qk_rope_head_dim
+ self.index_topk: int = config.index_topk
+ self.q_lora_rank: int = config.q_lora_rank
+
+ # Named to match checkpoint: wq_b, wk, k_norm
+ self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False)
+ self.wk = nn.Linear(self.hidden_size, self.head_dim, bias=False)
+ self.k_norm = nn.LayerNorm(self.head_dim, eps=1e-6)
+ # Named to match checkpoint: weights_proj
+ # In the reference, this is fp32; the HF FP8 checkpoint stores a bf16 tensor.
+ # Keeping it as a plain Linear prevents FP8 conversion (see `_keep_in_fp32_modules`).
+ self.weights_proj = nn.Linear(self.hidden_size, self.n_heads, bias=False)
+ self.softmax_scale = self.head_dim**-0.5
+
+ # Indexer maintains its own key cache (not in DynamicCache, which is sized for attention layers only)
+ self._cached_keys: torch.Tensor | None = None
+
+ @torch.no_grad()
+ def forward(
+ self,
+ hidden_states: torch.Tensor, # [B, S, hidden]
+ q_resid: torch.Tensor, # [B, S, q_lora_rank]
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: torch.Tensor | None,
+ use_cache: bool = False,
+ ) -> torch.LongTensor:
+ """
+ Computes top-k token indices for sparse attention (DSA).
+
+ This is the bf16 equivalent of the reference Indexer which uses `rotate_activation` (Hadamard transform)
+ and `fp8_index` (FP8 quantized scoring kernel). Since the Hadamard transform is orthogonal (dot products
+ are preserved: Hq·Hk = q·k), and FP8 quantization is a precision optimization, we skip both and compute
+ scores directly in bf16/fp32.
+
+ The scoring logic computes:
+ index_score[b,s,t] = Σ_h (weight[b,s,h] · softmax_scale · q[b,s,h,:] · k[b,t,:])
+
+ Args:
+ hidden_states: Input hidden states `[B, S, hidden_size]`.
+ q_resid: Query residual from `q_a_layernorm(q_a_proj(x))`, shape `[B, S, q_lora_rank]`.
+ position_embeddings: `(cos, sin)` from RotaryEmbedding.
+ attention_mask: Causal mask, broadcastable to `[B, S, T]`.
+ use_cache: Whether to store/update the indexer's own key cache for autoregressive decode.
+
+ Returns:
+ `torch.LongTensor`: Top-k token indices of shape `[B, S, topk]`.
+ """
+ batch_size, seq_len, _ = hidden_states.shape
+ cos, sin = position_embeddings
+
+ # === Queries ===
+ q = self.wq_b(q_resid) # [B, S, H*D]
+ q = q.view(batch_size, seq_len, self.n_heads, self.head_dim) # [B, S, H, D]
+ q_pe, q_nope = torch.split(q, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1)
+ q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=2) # [B, S, H, rope_D]
+ q = torch.cat([q_pe, q_nope], dim=-1) # [B, S, H, D]
+
+ # === Keys ===
+ k = self.k_norm(self.wk(hidden_states)) # [B, S, D]
+ k_pe, k_nope = torch.split(k, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1)
+ k_pe = apply_rotary_pos_emb(k_pe.unsqueeze(2), cos, sin, unsqueeze_dim=2).squeeze(2) # [B, S, rope_D]
+ k = torch.cat([k_pe, k_nope], dim=-1) # [B, S, D]
+
+ # === Key cache (managed by the indexer, not DynamicCache) ===
+ # Reset cache on prefill (new prompt) to avoid stale keys / batch-size mismatch
+ if seq_len > 1:
+ self._cached_keys = None
+
+ if use_cache:
+ if self._cached_keys is not None:
+ k_cached = torch.cat([self._cached_keys, k], dim=1) # [B, T, D]
+ else:
+ k_cached = k
+ self._cached_keys = k_cached
+ else:
+ k_cached = k
+
+ # === Scoring ===
+ # Reference: weights = weights_proj(x.float()) * n_heads^(-0.5)
+ # Reference: weights = weights.unsqueeze(-1) * q_scale * softmax_scale
+ # Reference: index_score = fp8_index(q_fp8, weights, k_cache, k_scale_cache)
+ #
+ # In bf16 mode (no FP8), q_scale = 1. The fp8_index kernel computes:
+ # score[b,s,t] = sum_h(weights[b,s,h] * dot(q[b,s,h,:], k[b,t,:]))
+ # where weights already absorbs n_heads^(-0.5) and softmax_scale.
+
+ # Don't force fp32 inputs here: the checkpoint stores `weights_proj.weight` in bf16.
+ # Use native dtype for matmul, then upcast the result for scoring stability.
+ weights = self.weights_proj(hidden_states).float() * (self.n_heads**-0.5) # [B, S, H]
+
+ # q·k^T per head: [B, S, H, D] @ [B, T, D]^T → [B, S, H, T]
+ scores = torch.einsum("bshd,btd->bsht", q.float(), k_cached.float()) * self.softmax_scale
+ scores = F.relu(scores)
+ # Weight per head and sum across heads → [B, S, T]
+ index_scores = torch.einsum("bsht,bsh->bst", scores, weights)
+
+ if attention_mask is not None:
+ index_scores = index_scores + attention_mask
+
+ total_len = index_scores.shape[-1]
+ topk = min(self.index_topk, total_len)
+ topk_indices = index_scores.topk(topk, dim=-1).indices # [B, S, topk]
+ return topk_indices
+
+
+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)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class GlmMoeDsaAttention(nn.Module):
+ """
+ Multi-head Latent Attention (MLA) with Dynamic Sparse Attention (DSA) indexer.
+
+ This follows the same architecture as DeepSeek V3.2's MLA:
+ - Query: x → q_a_proj → RMSNorm → q_b_proj → split(q_nope, q_pe) → RoPE(q_pe)
+ - KV: x → kv_a_proj → split(kv_compressed, k_pe) → RMSNorm(kv_compressed) → kv_b_proj
+ → RoPE(k_pe)
+ - Cache: fully expanded key_states [B, H, T, qk_head_dim] and value_states [B, H, T, v_head_dim]
+ - Indexer: selects top-k tokens via DSA, applied as an additive -inf mask on attention scores
+
+ **Caching strategy**: follows the DeepSeek V3 transformers convention of fully expanding K/V
+ before caching. This ensures compatibility with DynamicCache, StaticCache, flash attention,
+ and SDPA backends. The reference's compressed-cache decode path (which avoids the kv_b_proj
+ expansion at decode time) is a future optimization that would require a dedicated MLA cache class.
+
+ **FP8 compatibility**: all weight accesses use standard nn.Linear forward calls (never
+ raw `.weight` access), so FP8-quantized checkpoints work transparently.
+ """
+
+ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.attention_dropout = config.attention_dropout
+ self.num_heads = config.num_attention_heads
+
+ self.q_lora_rank = config.q_lora_rank
+ self.qk_rope_head_dim = config.qk_rope_head_dim
+ self.kv_lora_rank = config.kv_lora_rank
+ self.v_head_dim = config.v_head_dim
+ self.qk_nope_head_dim = config.qk_nope_head_dim
+ self.qk_head_dim = config.qk_head_dim
+
+ self.is_causal = True
+
+ # Query projection (with optional LoRA)
+ if self.q_lora_rank is None:
+ self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.qk_head_dim, bias=False)
+ else:
+ self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=config.attention_bias)
+ self.q_a_layernorm = GlmMoeDsaRMSNorm(config.q_lora_rank)
+ self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False)
+
+ # Key-Value projections (MLA compressed path)
+ self.kv_a_proj_with_mqa = nn.Linear(
+ config.hidden_size,
+ self.kv_lora_rank + self.qk_rope_head_dim,
+ bias=config.attention_bias,
+ )
+ self.kv_a_layernorm = GlmMoeDsaRMSNorm(self.kv_lora_rank)
+ self.kv_b_proj = nn.Linear(
+ self.kv_lora_rank,
+ self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
+ bias=False,
+ )
+
+ # Output projection
+ self.o_proj = nn.Linear(
+ self.num_heads * self.v_head_dim,
+ config.hidden_size,
+ bias=config.attention_bias,
+ )
+
+ self.scaling = self.qk_head_dim ** (-0.5)
+
+ self.indexer = GlmMoeDsaIndexer(config, layer_idx)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: torch.Tensor | None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+ batch_size, seq_length = hidden_states.shape[:-1]
+ cos, sin = position_embeddings
+
+ # ===== Query path =====
+ if self.q_lora_rank is None:
+ query_states = self.q_proj(hidden_states)
+ q_resid = None
+ else:
+ q_resid = self.q_a_layernorm(self.q_a_proj(hidden_states)) # [B, S, q_lora_rank]
+ query_states = self.q_b_proj(q_resid)
+ query_states = query_states.view(batch_size, seq_length, -1, self.qk_head_dim).transpose(1, 2)
+ # Split nope/rope, apply RoPE, recombine — layout: [B, H, S, D]
+ q_nope, q_pe = torch.split(query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+ q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # BHSD format
+
+ # ===== KV path =====
+ compressed_kv = self.kv_a_proj_with_mqa(hidden_states) # [B, S, kv_rank + rope_D]
+ k_compressed, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+ k_compressed = self.kv_a_layernorm(k_compressed) # [B, S, kv_rank]
+
+ # Expand KV through kv_b_proj
+ kv_expanded = self.kv_b_proj(k_compressed) # [B, S, H * (nope_D + v_D)]
+ kv_expanded = kv_expanded.view(batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim)
+ k_nope, value_states = torch.split(kv_expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+ k_nope = k_nope.transpose(1, 2) # [B, H, S, nope_D]
+ value_states = value_states.transpose(1, 2) # [B, H, S, v_D]
+
+ # RoPE on k_pe (single-head rope stream)
+ k_pe = k_pe.view(batch_size, 1, seq_length, self.qk_rope_head_dim) # [B, 1, S, rope_D]
+ k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # BHSD format
+ k_pe = k_pe.expand(-1, k_nope.shape[1], -1, -1) # [B, H, S, rope_D]
+
+ # Assemble full Q and K
+ query_states = torch.cat([q_nope, q_pe], dim=-1) # [B, H, S, qk_head_dim]
+ key_states = torch.cat([k_nope, k_pe], dim=-1) # [B, H, S, qk_head_dim]
+
+ # Cache update
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ # ===== Indexer (DSA sparse mask) =====
+ # attention_mask is [B, 1, S, T] (4D) for eager and (2D) otherwise but indexer works with [B, S, T] (3D)
+ indexer_mask = (
+ attention_mask[:, 0, :, :]
+ if attention_mask is not None and attention_mask.dim() == 4
+ else attention_mask.unsqueeze(1)
+ if attention_mask is not None
+ else None
+ )
+ topk_indices = self.indexer(
+ hidden_states,
+ q_resid,
+ position_embeddings,
+ indexer_mask,
+ use_cache=past_key_values is not None,
+ ) # [B, S, topk]
+
+ # Build combined DSA + causal mask: -inf everywhere except selected top-k positions
+ total_len = key_states.shape[2]
+ index_mask = torch.full(
+ (batch_size, seq_length, total_len),
+ float("-inf"),
+ device=hidden_states.device,
+ dtype=query_states.dtype,
+ )
+ index_mask.scatter_(-1, topk_indices, 0.0) # [B, S, T]
+ index_mask = index_mask.unsqueeze(1) # [B, 1, S, T]
+ if attention_mask is not None and attention_mask.dim() == 4:
+ causal_mask = attention_mask[..., :total_len]
+ combined_mask = index_mask + causal_mask
+ else:
+ combined_mask = (
+ attention_mask.masked_fill(index_mask == float("-inf"), float("-inf"))
+ if attention_mask is not None
+ else index_mask
+ )
+
+ # Flash attention head_dim padding (qk_head_dim != v_head_dim)
+ if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim:
+ value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim])
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ combined_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ indices=topk_indices, # flash_mla_with_kvcache
+ **kwargs,
+ )
+
+ if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim:
+ attn_output = attn_output[:, :, :, : self.v_head_dim]
+
+ attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class GlmMoeDsaMLP(nn.Module):
+ def __init__(self, config, intermediate_size=None):
+ super().__init__()
+ self.config = config
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size if intermediate_size is None else 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):
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+class GlmMoeDsaTopkRouter(nn.Module):
+ def __init__(self, config: GlmMoeDsaConfig):
+ super().__init__()
+ self.config = config
+ self.top_k = config.num_experts_per_tok
+ self.n_routed_experts = config.n_routed_experts
+ self.routed_scaling_factor = config.routed_scaling_factor
+ self.n_group = config.n_group
+ self.topk_group = config.topk_group
+ self.norm_topk_prob = config.norm_topk_prob
+
+ self.weight = nn.Parameter(torch.empty((self.n_routed_experts, config.hidden_size)))
+ self.register_buffer("e_score_correction_bias", torch.zeros((self.n_routed_experts), dtype=torch.float32))
+
+ def forward(self, hidden_states):
+ hidden_states = hidden_states.view(-1, self.config.hidden_size)
+ router_logits = F.linear(hidden_states.type(torch.float32), self.weight.type(torch.float32))
+ return router_logits
+
+
+@use_experts_implementation
+class GlmMoeDsaNaiveMoe(nn.Module):
+ """Collection of expert weights stored as 3D tensors."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.num_experts = config.num_local_experts
+ self.hidden_dim = config.hidden_size
+ self.intermediate_dim = config.moe_intermediate_size
+ self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim))
+ self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ top_k_index: torch.Tensor,
+ top_k_weights: torch.Tensor,
+ ) -> torch.Tensor:
+ final_hidden_states = torch.zeros_like(hidden_states)
+ with torch.no_grad():
+ expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
+ expert_mask = expert_mask.permute(2, 1, 0)
+ expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
+
+ for expert_idx in expert_hit:
+ expert_idx = expert_idx[0]
+ if expert_idx == self.num_experts:
+ continue
+ top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
+ current_state = hidden_states[token_idx]
+ gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
+ current_hidden_states = self.act_fn(gate) * up
+ current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
+ current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
+ final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
+
+ return final_hidden_states
+
+
+class GlmMoeDsaMoE(nn.Module):
+ """
+ A mixed expert module containing shared experts.
+ """
+
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.experts = GlmMoeDsaNaiveMoe(config)
+ self.gate = GlmMoeDsaTopkRouter(config)
+ self.shared_experts = GlmMoeDsaMLP(
+ config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts
+ )
+ self.n_routed_experts = config.n_routed_experts
+ self.n_group = config.n_group
+ self.topk_group = config.topk_group
+ self.norm_topk_prob = config.norm_topk_prob
+ self.routed_scaling_factor = config.routed_scaling_factor
+ self.top_k = config.num_experts_per_tok
+
+ def route_tokens_to_experts(self, router_logits):
+ router_logits = router_logits.sigmoid()
+ router_logits_for_choice = router_logits + self.gate.e_score_correction_bias
+ group_scores = (
+ router_logits_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group)
+ .topk(2, dim=-1)[0]
+ .sum(dim=-1)
+ )
+ group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
+ group_mask = torch.zeros_like(group_scores)
+ group_mask.scatter_(1, group_idx, 1)
+ score_mask = (
+ group_mask.unsqueeze(-1)
+ .expand(-1, self.n_group, self.n_routed_experts // self.n_group)
+ .reshape(-1, self.n_routed_experts)
+ )
+ scores_for_choice = router_logits_for_choice.masked_fill(~score_mask.bool(), 0.0)
+ topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]
+ topk_weights = router_logits.gather(1, topk_indices)
+ if self.norm_topk_prob:
+ denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20
+ topk_weights /= denominator
+ topk_weights = topk_weights * self.routed_scaling_factor
+ return topk_indices, topk_weights
+
+ def forward(self, hidden_states):
+ residuals = hidden_states
+ orig_shape = hidden_states.shape
+ router_logits = self.gate(hidden_states)
+ topk_indices, topk_weights = self.route_tokens_to_experts(router_logits)
+ hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
+ hidden_states = self.experts(hidden_states, topk_indices, topk_weights).view(*orig_shape)
+ hidden_states = hidden_states + self.shared_experts(residuals)
+ return hidden_states
+
+
+class GlmMoeDsaDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = GlmMoeDsaAttention(config, layer_idx)
+
+ if config.mlp_layer_types[layer_idx] == "sparse":
+ self.mlp = GlmMoeDsaMoE(config)
+ else:
+ self.mlp = GlmMoeDsaMLP(config)
+
+ self.input_layernorm = GlmMoeDsaRMSNorm(config.hidden_size, config.rms_norm_eps)
+ self.post_attention_layernorm = GlmMoeDsaRMSNorm(config.hidden_size, config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ 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
+ return hidden_states
+
+
+@auto_docstring
+class GlmMoeDsaPreTrainedModel(PreTrainedModel):
+ config: GlmMoeDsaConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["GlmMoeDsaDecoderLayer"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn = False # flash-mla kernels need a bit more work in the way we enable them!
+ _supports_sdpa = True
+ _supports_flex_attn = False
+
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": GlmMoeDsaDecoderLayer,
+ "attentions": GlmMoeDsaAttention,
+ }
+ _keep_in_fp32_modules_strict = ["e_score_correction_bias"]
+ _keys_to_ignore_on_load_unexpected = [r"model\.layers\.78.*"]
+ # NOTE: FP8 quantization uses `_keep_in_fp32_modules` (not `_strict`) to decide which modules to NOT convert.
+ # We must keep `indexer.weights_proj` as a plain Linear to match the checkpoint (no `weight_scale_inv`).
+ _keep_in_fp32_modules = ["indexer.weights_proj"]
+ _compatible_flash_implementations = ["kernels-community/flash-mla"]
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, GlmMoeDsaTopkRouter):
+ init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+ init.zeros_(module.e_score_correction_bias)
+ elif isinstance(module, GlmMoeDsaNaiveMoe):
+ init.normal_(module.gate_up_proj, mean=0.0, std=self.config.initializer_range)
+ init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range)
+
+
+class GlmMoeDsaRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: GlmMoeDsaConfig, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: GlmMoeDsaConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
+ head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+ dim = int(head_dim * partial_rotary_factor)
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@auto_docstring
+class GlmMoeDsaModel(GlmMoeDsaPreTrainedModel):
+ def __init__(self, config: GlmMoeDsaConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList(
+ [GlmMoeDsaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = GlmMoeDsaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = GlmMoeDsaRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_embeddings=position_embeddings,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class GlmMoeDsaForCausalLM(GlmMoeDsaPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+ _tp_plan = {"lm_head": "colwise_gather_output"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = GlmMoeDsaModel(config)
+ self.vocab_size = config.vocab_size
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> CausalLMOutputWithPast:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, GlmMoeDsaForCausalLM
+
+ >>> model = GlmMoeDsaForCausalLM.from_pretrained("meta-glm_moe_dsa/GlmMoeDsa-2-7b-hf")
+ >>> tokenizer = AutoTokenizer.from_pretrained("meta-glm_moe_dsa/GlmMoeDsa-2-7b-hf")
+
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+ ```"""
+ outputs: BaseModelOutputWithPast = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = ["GlmMoeDsaPreTrainedModel", "GlmMoeDsaModel", "GlmMoeDsaForCausalLM"]
diff --git a/third_party/transformers/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py b/third_party/transformers/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py
new file mode 100644
index 0000000000000000000000000000000000000000..fcae77cb25629792ab57bf26a300b35c3b3ca95a
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_moe_dsa/modular_glm_moe_dsa.py
@@ -0,0 +1,490 @@
+# Copyright 2026 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from collections.abc import Callable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub.dataclasses import strict
+
+from ...cache_utils import Cache
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
+from ...models.llama.modeling_llama import rotate_half
+from ...processing_utils import Unpack
+from ...utils import auto_docstring, logging
+from ...utils.generic import is_flash_attention_requested
+from ..glm4_moe.modeling_glm4_moe import (
+ Glm4MoeForCausalLM,
+ Glm4MoeModel,
+ Glm4MoePreTrainedModel,
+ Glm4MoeRMSNorm,
+)
+from ..glm4_moe_lite.configuration_glm4_moe_lite import Glm4MoeLiteConfig
+from ..glm4_moe_lite.modeling_glm4_moe_lite import (
+ Glm4MoeLiteDecoderLayer,
+ eager_attention_forward,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+def apply_rotary_pos_emb(
+ x: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+ unsqueeze_dim: int = 1,
+) -> torch.Tensor:
+ """
+ Applies Rotary Position Embedding to a single tensor.
+
+ This is the transformers equivalent of DeepSeek V3.2's `apply_rotary_emb(x, freqs_cis, interleaved)`.
+ Instead of using complex-number `freqs_cis`, we use pre-split `(cos, sin)` tensors from RotaryEmbedding.
+
+ Args:
+ x (`torch.Tensor`): Input tensor of shape `[..., head_dim]`.
+ cos (`torch.Tensor`): Cosine part from RotaryEmbedding, shape `[batch, seq_len, head_dim]`.
+ sin (`torch.Tensor`): Sine part from RotaryEmbedding, shape `[batch, seq_len, head_dim]`.
+ unsqueeze_dim (`int`): Dimension along which to unsqueeze cos/sin for broadcasting.
+ Use `1` when x is `[B, H, S, D]` (BHSD) and `2` when x is `[B, S, H, D]` (BSHD).
+
+ Returns:
+ `torch.Tensor`: Tensor with rotary embeddings applied, same shape as input.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+
+ # Split-half (NeoX/Llama style): (x[:d/2], x[d/2:])
+ # This matches llama's apply_rotary_pos_emb logic.
+ x_rotated = (x * cos) + (rotate_half(x) * sin)
+ return x_rotated
+
+
+@auto_docstring(checkpoint="zai-org/GLM-5")
+@strict
+class GlmMoeDsaConfig(Glm4MoeLiteConfig):
+ r"""
+ n_group (`int`, *optional*, defaults to 1):
+ Number of groups for routed experts.
+ mlp_layer_types (`list`, *optional*):
+ MLP type pattern for each layer (`"dense"` or `"sparse"`). Defaults to 3 dense + rest sparse.
+ index_topk (`int`, *optional*, defaults to 2048):
+ Number of top tokens selected by the indexer for sparse attention.
+ index_head_dim (`int`, *optional*, defaults to 128):
+ Head dimension for the indexer projections (DSA).
+ index_n_heads (`int | None`, *optional*, defaults to 32):
+ Number of heads for the indexer projections (DSA).
+
+ ```python
+ >>> from transformers import GlmMoeDsaConfig, GlmMoeDsaModel
+
+ >>> # Initializing a GLM-MoE-DSA configuration
+ >>> configuration = GlmMoeDsaConfig()
+
+ >>> # Initializing a model from the configuration
+ >>> model = GlmMoeDsaModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_b_proj": "colwise",
+ "layers.*.self_attn.kv_a_proj_with_mqa": "mla_kv_a_proj",
+ "layers.*.self_attn.kv_b_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.experts.gate_up_proj": "packed_colwise",
+ "layers.*.mlp.experts.down_proj": "rowwise",
+ "layers.*.mlp.experts": "moe_tp_experts",
+ "layers.*.mlp.shared_experts.gate_proj": "colwise",
+ "layers.*.mlp.shared_experts.up_proj": "colwise",
+ "layers.*.mlp.shared_experts.down_proj": "rowwise",
+ "layers.*.mlp.gate_proj": "colwise",
+ "layers.*.mlp.up_proj": "colwise",
+ "layers.*.mlp.down_proj": "rowwise",
+ }
+
+ hidden_size: int = 6144
+ intermediate_size: int = 12288
+ moe_intermediate_size: int = 2048
+ num_hidden_layers: int = 78
+ num_attention_heads: int = 64
+ num_key_value_heads: int = 64
+ n_routed_experts: int = 256
+ routed_scaling_factor: float = 2.5
+ q_lora_rank: int = 2048
+ num_experts_per_tok: int = 8
+ index_topk: int = 2048
+ index_head_dim: int = 128
+ index_n_heads: int = 32
+ pretraining_tp = AttributeError()
+ rope_interleave = AttributeError()
+
+ def __post_init__(self, **kwargs):
+ self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+
+ # MLP layer types: first 3 dense, rest sparse
+ if self.mlp_layer_types is None:
+ self.mlp_layer_types = ["dense"] * min(3, self.num_hidden_layers) + ["sparse"] * (
+ self.num_hidden_layers - 3
+ )
+ PreTrainedConfig.__post_init__(self, **kwargs)
+
+
+class GlmMoeDsaRMSNorm(Glm4MoeRMSNorm):
+ pass
+
+
+class GlmMoeDsaIndexer(nn.Module):
+ """
+ Dynamic Sparse Attention (DSA) indexer for selecting top-k tokens.
+
+ The Indexer has its own lightweight projections (wq_b, wk) separate from the
+ main MLA attention. It uses non-interleaved (NeoX/Llama) RoPE, unlike the main attention
+ which uses interleaved RoPE.
+
+ **Cache strategy**: The Indexer manages its own key cache (`_cached_keys`) separately
+ from the DynamicCache used by MLA attention, since DynamicCache is sized for exactly
+ `num_hidden_layers` attention layers. Keys are concatenated along the sequence dimension
+ during autoregressive decode.
+ """
+
+ def __init__(self, config: "GlmMoeDsaConfig", layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+
+ self.hidden_size: int = config.hidden_size
+ self.n_heads: int = config.index_n_heads
+ self.head_dim: int = config.index_head_dim
+ self.qk_rope_head_dim: int = config.qk_rope_head_dim
+ self.index_topk: int = config.index_topk
+ self.q_lora_rank: int = config.q_lora_rank
+
+ # Named to match checkpoint: wq_b, wk, k_norm
+ self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.head_dim, bias=False)
+ self.wk = nn.Linear(self.hidden_size, self.head_dim, bias=False)
+ self.k_norm = nn.LayerNorm(self.head_dim, eps=1e-6)
+ # Named to match checkpoint: weights_proj
+ # In the reference, this is fp32; the HF FP8 checkpoint stores a bf16 tensor.
+ # Keeping it as a plain Linear prevents FP8 conversion (see `_keep_in_fp32_modules`).
+ self.weights_proj = nn.Linear(self.hidden_size, self.n_heads, bias=False)
+ self.softmax_scale = self.head_dim**-0.5
+
+ # Indexer maintains its own key cache (not in DynamicCache, which is sized for attention layers only)
+ self._cached_keys: torch.Tensor | None = None
+
+ @torch.no_grad()
+ def forward(
+ self,
+ hidden_states: torch.Tensor, # [B, S, hidden]
+ q_resid: torch.Tensor, # [B, S, q_lora_rank]
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: torch.Tensor | None,
+ use_cache: bool = False,
+ ) -> torch.LongTensor:
+ """
+ Computes top-k token indices for sparse attention (DSA).
+
+ This is the bf16 equivalent of the reference Indexer which uses `rotate_activation` (Hadamard transform)
+ and `fp8_index` (FP8 quantized scoring kernel). Since the Hadamard transform is orthogonal (dot products
+ are preserved: Hq·Hk = q·k), and FP8 quantization is a precision optimization, we skip both and compute
+ scores directly in bf16/fp32.
+
+ The scoring logic computes:
+ index_score[b,s,t] = Σ_h (weight[b,s,h] · softmax_scale · q[b,s,h,:] · k[b,t,:])
+
+ Args:
+ hidden_states: Input hidden states `[B, S, hidden_size]`.
+ q_resid: Query residual from `q_a_layernorm(q_a_proj(x))`, shape `[B, S, q_lora_rank]`.
+ position_embeddings: `(cos, sin)` from RotaryEmbedding.
+ attention_mask: Causal mask, broadcastable to `[B, S, T]`.
+ use_cache: Whether to store/update the indexer's own key cache for autoregressive decode.
+
+ Returns:
+ `torch.LongTensor`: Top-k token indices of shape `[B, S, topk]`.
+ """
+ batch_size, seq_len, _ = hidden_states.shape
+ cos, sin = position_embeddings
+
+ # === Queries ===
+ q = self.wq_b(q_resid) # [B, S, H*D]
+ q = q.view(batch_size, seq_len, self.n_heads, self.head_dim) # [B, S, H, D]
+ q_pe, q_nope = torch.split(q, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1)
+ q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=2) # [B, S, H, rope_D]
+ q = torch.cat([q_pe, q_nope], dim=-1) # [B, S, H, D]
+
+ # === Keys ===
+ k = self.k_norm(self.wk(hidden_states)) # [B, S, D]
+ k_pe, k_nope = torch.split(k, [self.qk_rope_head_dim, self.head_dim - self.qk_rope_head_dim], dim=-1)
+ k_pe = apply_rotary_pos_emb(k_pe.unsqueeze(2), cos, sin, unsqueeze_dim=2).squeeze(2) # [B, S, rope_D]
+ k = torch.cat([k_pe, k_nope], dim=-1) # [B, S, D]
+
+ # === Key cache (managed by the indexer, not DynamicCache) ===
+ # Reset cache on prefill (new prompt) to avoid stale keys / batch-size mismatch
+ if seq_len > 1:
+ self._cached_keys = None
+
+ if use_cache:
+ if self._cached_keys is not None:
+ k_cached = torch.cat([self._cached_keys, k], dim=1) # [B, T, D]
+ else:
+ k_cached = k
+ self._cached_keys = k_cached
+ else:
+ k_cached = k
+
+ # === Scoring ===
+ # Reference: weights = weights_proj(x.float()) * n_heads^(-0.5)
+ # Reference: weights = weights.unsqueeze(-1) * q_scale * softmax_scale
+ # Reference: index_score = fp8_index(q_fp8, weights, k_cache, k_scale_cache)
+ #
+ # In bf16 mode (no FP8), q_scale = 1. The fp8_index kernel computes:
+ # score[b,s,t] = sum_h(weights[b,s,h] * dot(q[b,s,h,:], k[b,t,:]))
+ # where weights already absorbs n_heads^(-0.5) and softmax_scale.
+
+ # Don't force fp32 inputs here: the checkpoint stores `weights_proj.weight` in bf16.
+ # Use native dtype for matmul, then upcast the result for scoring stability.
+ weights = self.weights_proj(hidden_states).float() * (self.n_heads**-0.5) # [B, S, H]
+
+ # q·k^T per head: [B, S, H, D] @ [B, T, D]^T → [B, S, H, T]
+ scores = torch.einsum("bshd,btd->bsht", q.float(), k_cached.float()) * self.softmax_scale
+ scores = F.relu(scores)
+ # Weight per head and sum across heads → [B, S, T]
+ index_scores = torch.einsum("bsht,bsh->bst", scores, weights)
+
+ if attention_mask is not None:
+ index_scores = index_scores + attention_mask
+
+ total_len = index_scores.shape[-1]
+ topk = min(self.index_topk, total_len)
+ topk_indices = index_scores.topk(topk, dim=-1).indices # [B, S, topk]
+ return topk_indices
+
+
+class GlmMoeDsaAttention(nn.Module):
+ """
+ Multi-head Latent Attention (MLA) with Dynamic Sparse Attention (DSA) indexer.
+
+ This follows the same architecture as DeepSeek V3.2's MLA:
+ - Query: x → q_a_proj → RMSNorm → q_b_proj → split(q_nope, q_pe) → RoPE(q_pe)
+ - KV: x → kv_a_proj → split(kv_compressed, k_pe) → RMSNorm(kv_compressed) → kv_b_proj
+ → RoPE(k_pe)
+ - Cache: fully expanded key_states [B, H, T, qk_head_dim] and value_states [B, H, T, v_head_dim]
+ - Indexer: selects top-k tokens via DSA, applied as an additive -inf mask on attention scores
+
+ **Caching strategy**: follows the DeepSeek V3 transformers convention of fully expanding K/V
+ before caching. This ensures compatibility with DynamicCache, StaticCache, flash attention,
+ and SDPA backends. The reference's compressed-cache decode path (which avoids the kv_b_proj
+ expansion at decode time) is a future optimization that would require a dedicated MLA cache class.
+
+ **FP8 compatibility**: all weight accesses use standard nn.Linear forward calls (never
+ raw `.weight` access), so FP8-quantized checkpoints work transparently.
+ """
+
+ def __init__(self, config: GlmMoeDsaConfig, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.attention_dropout = config.attention_dropout
+ self.num_heads = config.num_attention_heads
+
+ self.q_lora_rank = config.q_lora_rank
+ self.qk_rope_head_dim = config.qk_rope_head_dim
+ self.kv_lora_rank = config.kv_lora_rank
+ self.v_head_dim = config.v_head_dim
+ self.qk_nope_head_dim = config.qk_nope_head_dim
+ self.qk_head_dim = config.qk_head_dim
+
+ self.is_causal = True
+
+ # Query projection (with optional LoRA)
+ if self.q_lora_rank is None:
+ self.q_proj = nn.Linear(config.hidden_size, self.num_heads * self.qk_head_dim, bias=False)
+ else:
+ self.q_a_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=config.attention_bias)
+ self.q_a_layernorm = GlmMoeDsaRMSNorm(config.q_lora_rank)
+ self.q_b_proj = nn.Linear(config.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False)
+
+ # Key-Value projections (MLA compressed path)
+ self.kv_a_proj_with_mqa = nn.Linear(
+ config.hidden_size,
+ self.kv_lora_rank + self.qk_rope_head_dim,
+ bias=config.attention_bias,
+ )
+ self.kv_a_layernorm = GlmMoeDsaRMSNorm(self.kv_lora_rank)
+ self.kv_b_proj = nn.Linear(
+ self.kv_lora_rank,
+ self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
+ bias=False,
+ )
+
+ # Output projection
+ self.o_proj = nn.Linear(
+ self.num_heads * self.v_head_dim,
+ config.hidden_size,
+ bias=config.attention_bias,
+ )
+
+ self.scaling = self.qk_head_dim ** (-0.5)
+
+ self.indexer = GlmMoeDsaIndexer(config, layer_idx)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: torch.Tensor | None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+ batch_size, seq_length = hidden_states.shape[:-1]
+ cos, sin = position_embeddings
+
+ # ===== Query path =====
+ if self.q_lora_rank is None:
+ query_states = self.q_proj(hidden_states)
+ q_resid = None
+ else:
+ q_resid = self.q_a_layernorm(self.q_a_proj(hidden_states)) # [B, S, q_lora_rank]
+ query_states = self.q_b_proj(q_resid)
+ query_states = query_states.view(batch_size, seq_length, -1, self.qk_head_dim).transpose(1, 2)
+ # Split nope/rope, apply RoPE, recombine — layout: [B, H, S, D]
+ q_nope, q_pe = torch.split(query_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+ q_pe = apply_rotary_pos_emb(q_pe, cos, sin, unsqueeze_dim=1) # BHSD format
+
+ # ===== KV path =====
+ compressed_kv = self.kv_a_proj_with_mqa(hidden_states) # [B, S, kv_rank + rope_D]
+ k_compressed, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+ k_compressed = self.kv_a_layernorm(k_compressed) # [B, S, kv_rank]
+
+ # Expand KV through kv_b_proj
+ kv_expanded = self.kv_b_proj(k_compressed) # [B, S, H * (nope_D + v_D)]
+ kv_expanded = kv_expanded.view(batch_size, seq_length, -1, self.qk_nope_head_dim + self.v_head_dim)
+ k_nope, value_states = torch.split(kv_expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+ k_nope = k_nope.transpose(1, 2) # [B, H, S, nope_D]
+ value_states = value_states.transpose(1, 2) # [B, H, S, v_D]
+
+ # RoPE on k_pe (single-head rope stream)
+ k_pe = k_pe.view(batch_size, 1, seq_length, self.qk_rope_head_dim) # [B, 1, S, rope_D]
+ k_pe = apply_rotary_pos_emb(k_pe, cos, sin, unsqueeze_dim=1) # BHSD format
+ k_pe = k_pe.expand(-1, k_nope.shape[1], -1, -1) # [B, H, S, rope_D]
+
+ # Assemble full Q and K
+ query_states = torch.cat([q_nope, q_pe], dim=-1) # [B, H, S, qk_head_dim]
+ key_states = torch.cat([k_nope, k_pe], dim=-1) # [B, H, S, qk_head_dim]
+
+ # Cache update
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ # ===== Indexer (DSA sparse mask) =====
+ # attention_mask is [B, 1, S, T] (4D) for eager and (2D) otherwise but indexer works with [B, S, T] (3D)
+ indexer_mask = (
+ attention_mask[:, 0, :, :]
+ if attention_mask is not None and attention_mask.dim() == 4
+ else attention_mask.unsqueeze(1)
+ if attention_mask is not None
+ else None
+ )
+ topk_indices = self.indexer(
+ hidden_states,
+ q_resid,
+ position_embeddings,
+ indexer_mask,
+ use_cache=past_key_values is not None,
+ ) # [B, S, topk]
+
+ # Build combined DSA + causal mask: -inf everywhere except selected top-k positions
+ total_len = key_states.shape[2]
+ index_mask = torch.full(
+ (batch_size, seq_length, total_len),
+ float("-inf"),
+ device=hidden_states.device,
+ dtype=query_states.dtype,
+ )
+ index_mask.scatter_(-1, topk_indices, 0.0) # [B, S, T]
+ index_mask = index_mask.unsqueeze(1) # [B, 1, S, T]
+ if attention_mask is not None and attention_mask.dim() == 4:
+ causal_mask = attention_mask[..., :total_len]
+ combined_mask = index_mask + causal_mask
+ else:
+ combined_mask = (
+ attention_mask.masked_fill(index_mask == float("-inf"), float("-inf"))
+ if attention_mask is not None
+ else index_mask
+ )
+
+ # Flash attention head_dim padding (qk_head_dim != v_head_dim)
+ if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim:
+ value_states = F.pad(value_states, [0, self.qk_head_dim - self.v_head_dim])
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ combined_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ indices=topk_indices, # flash_mla_with_kvcache
+ **kwargs,
+ )
+
+ if is_flash_attention_requested(self.config) and self.qk_head_dim != self.v_head_dim:
+ attn_output = attn_output[:, :, :, : self.v_head_dim]
+
+ attn_output = attn_output.reshape(batch_size, seq_length, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class GlmMoeDsaDecoderLayer(Glm4MoeLiteDecoderLayer):
+ pass
+
+
+class GlmMoeDsaPreTrainedModel(Glm4MoePreTrainedModel):
+ # NOTE: FP8 quantization uses `_keep_in_fp32_modules` (not `_strict`) to decide which modules to NOT convert.
+ # We must keep `indexer.weights_proj` as a plain Linear to match the checkpoint (no `weight_scale_inv`).
+ _keep_in_fp32_modules = ["indexer.weights_proj"]
+ _keep_in_fp32_modules_strict = ["e_score_correction_bias"]
+ _keys_to_ignore_on_load_unexpected = [r"model\.layers\.78.*"]
+ _supports_flash_attn = False # flash-mla kernels need a bit more work in the way we enable them!
+ _supports_sdpa = True
+ _supports_flex_attn = False
+ _compatible_flash_implementations = ["kernels-community/flash-mla"]
+
+
+class GlmMoeDsaModel(Glm4MoeModel):
+ pass
+
+
+class GlmMoeDsaForCausalLM(Glm4MoeForCausalLM):
+ pass
+
+
+__all__ = [
+ "GlmMoeDsaConfig",
+ "GlmMoeDsaPreTrainedModel",
+ "GlmMoeDsaModel",
+ "GlmMoeDsaForCausalLM",
+]
diff --git a/third_party/transformers/src/transformers/models/glm_ocr/__init__.py b/third_party/transformers/src/transformers/models/glm_ocr/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..528a9155b68635cd3acc5b9a40e31a9fd1f9ae70
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_ocr/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2026 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_glm_ocr import *
+ from .modeling_glm_ocr import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/glm_ocr/configuration_glm_ocr.py b/third_party/transformers/src/transformers/models/glm_ocr/configuration_glm_ocr.py
new file mode 100644
index 0000000000000000000000000000000000000000..d08710c5e8b5acd0f548c06258e8abd402bc85e7
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_ocr/configuration_glm_ocr.py
@@ -0,0 +1,185 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/glm_ocr/modular_glm_ocr.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_glm_ocr.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2026 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="zai-org/GLM-OCR")
+@strict
+class GlmOcrVisionConfig(PreTrainedConfig):
+ r"""
+ out_hidden_size (`int`, *optional*, defaults to 4096):
+ The output hidden size of the vision model.
+
+ Example:
+
+ ```python
+ >>> from transformers import GlmOcrVisionConfig, GlmOcrVisionModel
+
+ >>> # Initializing a GlmOcrVisionConfig GLM-4.1V-9B style configuration
+ >>> configuration = GlmOcrVisionConfig()
+
+ >>> # Initializing a model (with random weights) from the GLM-4.1V-9B configuration
+ >>> model = GlmOcrVisionModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "glm_ocr_vision"
+ base_config_key = "vision_config"
+
+ depth: int = 24
+ hidden_size: int = 1024
+ hidden_act: str = "silu"
+ attention_bias: bool = True
+ attention_dropout: float | int = 0.0
+ num_heads: int = 16
+ in_channels: int = 3
+ image_size: int | list[int] | tuple[int, int] = 336
+ patch_size: int | list[int] | tuple[int, int] = 14
+ rms_norm_eps: float = 1e-05
+ spatial_merge_size: int = 2
+ temporal_patch_size: int | list[int] | tuple[int, int] = 2
+ out_hidden_size: int = 1536
+ intermediate_size: int = 4096
+ initializer_range: float = 0.02
+
+
+@auto_docstring(checkpoint="zai-org/GLM-OCR")
+@strict
+class GlmOcrTextConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import GlmOcrTextModel, GlmOcrConfig
+
+ >>> # Initializing a GLM-OCR style configuration
+ >>> configuration = GlmOcrConfig()
+
+ >>> # Initializing a model from the GLM-OCR style configuration
+ >>> model = GlmOcrTextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "glm_ocr_text"
+ base_config_key = "text_config"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ # Default tensor parallel plan for base model `GlmOcr`
+ base_model_tp_plan = {
+ "layers.*.self_attn.q_proj": "colwise",
+ "layers.*.self_attn.k_proj": "colwise",
+ "layers.*.self_attn.v_proj": "colwise",
+ "layers.*.self_attn.o_proj": "rowwise",
+ "layers.*.mlp.gate_up_proj": "colwise_gather_output", # we need to replicate here due to the `chunk` operation
+ "layers.*.mlp.down_proj": "rowwise_split_input", # input is replicated due to the `chunk` operation
+ }
+ base_model_pp_plan = {
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "norm": (["hidden_states"], ["hidden_states"]),
+ }
+ ignore_keys_at_rope_validation = {"mrope_section"}
+
+ vocab_size: int = 59392
+ hidden_size: int = 1024
+ intermediate_size: int = 4096
+ num_hidden_layers: int = 16
+ num_attention_heads: int = 16
+ num_key_value_heads: int = 8
+ hidden_act: str = "silu"
+ max_position_embeddings: int = 131072
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-05
+ use_cache: bool = True
+ attention_dropout: float | int = 0.0
+ rope_parameters: RopeParameters | dict | None = None
+ pad_token_id: int | None = None
+
+ def __post_init__(self, **kwargs):
+ if self.num_key_value_heads is None:
+ self.num_key_value_heads = self.num_attention_heads
+
+ super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="zai-org/GLM-OCR")
+@strict
+class GlmOcrConfig(PreTrainedConfig):
+ r"""
+ image_start_token_id (`int`, *optional*, defaults to 59256):
+ The image start token index to encode the start of image.
+ image_end_token_id (`int`, *optional*, defaults to 59257):
+ The image end token index to encode the end of image.
+ video_start_token_id (`int`, *optional*, defaults to 59258):
+ The video start token index to encode the start of video.
+ video_end_token_id (`int`, *optional*, defaults to 59259):
+ The video end token index to encode the end of video.
+
+ ```python
+ >>> from transformers import GlmOcrForConditionalGeneration, GlmOcrConfig
+
+ >>> # Initializing a GLM-OCR style configuration
+ >>> configuration = GlmOcrConfig()
+
+ >>> # Initializing a model from the GLM-OCR style configuration
+ >>> model = GlmOcrForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "glm_ocr"
+ sub_configs = {"vision_config": GlmOcrVisionConfig, "text_config": GlmOcrTextConfig}
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ text_config: dict | PreTrainedConfig | None = None
+ vision_config: dict | PreTrainedConfig | None = None
+
+ image_token_id: int = 59280
+ video_token_id: int = 59281
+ image_start_token_id: int = 59256
+ image_end_token_id: int = 59257
+ video_start_token_id: int = 59258
+ video_end_token_id: int = 59259
+ tie_word_embeddings: bool = False
+
+ def __post_init__(self, **kwargs):
+ if isinstance(self.vision_config, dict):
+ self.vision_config = self.sub_configs["vision_config"](**self.vision_config)
+ elif self.vision_config is None:
+ self.vision_config = self.sub_configs["vision_config"](**kwargs)
+
+ if isinstance(self.text_config, dict):
+ self.text_config = self.sub_configs["text_config"](**self.text_config)
+ elif self.text_config is None:
+ self.text_config = self.sub_configs["text_config"](**kwargs)
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["GlmOcrConfig", "GlmOcrTextConfig", "GlmOcrVisionConfig"]
diff --git a/third_party/transformers/src/transformers/models/glm_ocr/modeling_glm_ocr.py b/third_party/transformers/src/transformers/models/glm_ocr/modeling_glm_ocr.py
new file mode 100644
index 0000000000000000000000000000000000000000..30703d81c8c10d8404bbf1e7864769c01ffe19fe
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_ocr/modeling_glm_ocr.py
@@ -0,0 +1,1650 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/glm_ocr/modular_glm_ocr.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_glm_ocr.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2026 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+import itertools
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import LayerNorm
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, BaseModelOutputWithPooling, ModelOutput
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, torch_compilable_check
+from ...utils.generic import is_flash_attention_requested, maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_glm_ocr import GlmOcrConfig, GlmOcrTextConfig, GlmOcrVisionConfig
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class GlmOcrRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ GlmOcrRMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class GlmOcrVisionMlp(nn.Module):
+ def __init__(self, config, bias: bool = True):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.intermediate_size = config.intermediate_size
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=bias)
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias)
+ self.act_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_state):
+ return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
+
+
+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)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+def rotate_half_llm(x):
+ """Rotates half the hidden dims of the input."""
+ x1 = x[..., 0::2]
+ x2 = x[..., 1::2]
+ return torch.stack((-x2, x1), dim=-1).flatten(-2)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+
+ # Interleave them instead of usual shape
+ cos = cos[..., : cos.shape[-1] // 2].repeat_interleave(2, dim=-1)
+ sin = sin[..., : sin.shape[-1] // 2].repeat_interleave(2, dim=-1)
+
+ # Keep half or full tensor for later concatenation
+ rotary_dim = cos.shape[-1]
+ q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+ k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+ # Apply rotary embeddings on the first half or full tensor
+ q_embed = (q_rot * cos) + (rotate_half_llm(q_rot) * sin)
+ k_embed = (k_rot * cos) + (rotate_half_llm(k_rot) * sin)
+
+ # Concatenate back to full shape
+ q_embed = torch.cat([q_embed, q_pass], dim=-1)
+ k_embed = torch.cat([k_embed, k_pass], dim=-1)
+ return q_embed, k_embed
+
+
+class GlmOcrTextAttention(nn.Module):
+ """
+ Multi-headed attention from 'Attention Is All You Need' paper.
+ and "Generating Long Sequences with Sparse Transformers".
+ """
+
+ def __init__(self, config: GlmOcrTextConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+
+ self.hidden_size = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_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.is_causal = True
+ self.attention_dropout = config.attention_dropout
+ self.rope_parameters = config.rope_parameters
+ self.scaling = self.head_dim**-0.5
+ 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)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+ 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, -1, self.head_dim).transpose(1, 2)
+ key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+ value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class GlmOcrTextMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+
+ self.config = config
+ self.gate_up_proj = nn.Linear(config.hidden_size, 2 * config.intermediate_size, bias=False)
+ self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
+ self.activation_fn = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
+ up_states = self.gate_up_proj(hidden_states)
+
+ gate, up_states = up_states.chunk(2, dim=-1)
+ up_states = up_states * self.activation_fn(gate)
+
+ return self.down_proj(up_states)
+
+
+class GlmOcrTextDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: GlmOcrTextConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = GlmOcrTextAttention(config, layer_idx)
+ self.mlp = GlmOcrTextMLP(config)
+ self.input_layernorm = GlmOcrRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = GlmOcrRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_self_attn_layernorm = GlmOcrRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_mlp_layernorm = GlmOcrRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ @auto_docstring
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ **kwargs,
+ ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = self.post_self_attn_layernorm(hidden_states)
+ 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 = self.post_mlp_layernorm(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+class GlmOcrVisionRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, dim: int, theta: float = 10000.0) -> None:
+ super().__init__()
+ self.dim = dim
+ self.theta = theta
+ inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+ def forward(self, seqlen: int) -> torch.Tensor:
+ seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
+ freqs = torch.outer(seq, self.inv_freq)
+ return freqs
+
+
+@auto_docstring
+class GlmOcrPreTrainedModel(PreTrainedModel):
+ config: GlmOcrConfig
+ base_model_prefix = "model"
+ input_modalities = ("image", "video", "text")
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["GlmOcrTextDecoderLayer", "GlmOcrVisionBlock"]
+ _skip_keys_device_placement = "past_key_values"
+ _supports_flash_attn = True
+ _supports_sdpa = True
+
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+ _keys_to_ignore_on_load_unexpected = [r"model\.language_model\.layers\.16.*"]
+
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, GlmOcrVisionRotaryEmbedding):
+ inv_freq = 1.0 / (module.theta ** (torch.arange(0, module.dim, 2, dtype=torch.float) / module.dim))
+ init.copy_(module.inv_freq, inv_freq)
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for Llava outputs, with hidden states and attentions.
+ """
+)
+class GlmOcrModelOutputWithPast(ModelOutput):
+ r"""
+ past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ last_hidden_state: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ rope_deltas: torch.LongTensor | None = None
+
+
+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)
+
+
+def apply_rotary_pos_emb_vision(
+ q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> tuple[torch.Tensor, torch.Tensor]:
+ orig_q_dtype = q.dtype
+ orig_k_dtype = k.dtype
+ q, k = q.float(), k.float()
+ cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ q_embed = q_embed.to(orig_q_dtype)
+ k_embed = k_embed.to(orig_k_dtype)
+ return q_embed, k_embed
+
+
+class GlmOcrVisionAttention(nn.Module):
+ def __init__(self, config: GlmOcrVisionConfig) -> None:
+ super().__init__()
+ self.dim = config.hidden_size
+ self.num_heads = config.num_heads
+ self.head_dim = self.dim // self.num_heads
+ self.num_key_value_groups = 1 # needed for eager attention
+ self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.attention_bias)
+ self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)
+ self.scaling = self.head_dim**-0.5
+ self.config = config
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = False
+ self.q_norm = GlmOcrRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+ self.k_norm = GlmOcrRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: torch.Tensor | None = None,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ query_states, key_states, value_states = (
+ self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+ )
+
+ query_states = self.q_norm(query_states)
+ key_states = self.k_norm(key_states)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)
+ query_states = query_states.transpose(0, 1).unsqueeze(0)
+ key_states = key_states.transpose(0, 1).unsqueeze(0)
+ value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ if is_flash_attention_requested(self.config):
+ # Flash Attention: Use cu_seqlens for variable length attention
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+ attn_output, _ = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ cu_seq_lens_q=cu_seqlens,
+ cu_seq_lens_k=cu_seqlens,
+ max_length_q=max_seqlen,
+ max_length_k=max_seqlen,
+ is_causal=False,
+ **kwargs,
+ )
+ else:
+ # Other implementations: Process each chunk separately
+ lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+ splits = [
+ torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
+ ]
+
+ attn_outputs = [
+ attention_interface(
+ self,
+ q,
+ k,
+ v,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ is_causal=False,
+ **kwargs,
+ )[0]
+ for q, k, v in zip(*splits)
+ ]
+ attn_output = torch.cat(attn_outputs, dim=1)
+
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+class GlmOcrVisionBlock(GradientCheckpointingLayer):
+ def __init__(self, config) -> None:
+ super().__init__()
+ self.norm1 = GlmOcrRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.norm2 = GlmOcrRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.attn = GlmOcrVisionAttention(config)
+ self.mlp = GlmOcrVisionMlp(config, bias=config.attention_bias)
+
+ @auto_docstring
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: torch.Tensor | None = None,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ r"""
+ cu_seqlens (`torch.Tensor`):
+ Cumulative sequence lengths used for packed variable-length attention in Flash Attention kernels.
+ rotary_pos_emb (`torch.Tensor`, *optional*):
+ Precomputed rotary positional embeddings applied to the vision attention query/key states.
+ """
+ hidden_states = hidden_states + self.attn(
+ self.norm1(hidden_states),
+ cu_seqlens=cu_seqlens,
+ rotary_pos_emb=rotary_pos_emb,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = hidden_states + self.mlp(self.norm2(hidden_states))
+ return hidden_states
+
+
+class GlmOcrVisionPatchMerger(nn.Module):
+ def __init__(self, dim: int, context_dim: int, hidden_act: str, bias: bool = False) -> None:
+ super().__init__()
+ self.proj = nn.Linear(dim, dim, bias=bias)
+ self.post_projection_norm = LayerNorm(dim)
+ self.gate_proj = nn.Linear(dim, context_dim, bias=bias)
+ self.up_proj = nn.Linear(dim, context_dim, bias=bias)
+ self.down_proj = nn.Linear(context_dim, dim, bias=bias)
+ self.act1 = nn.GELU()
+ self.act_fn = ACT2FN[hidden_act]
+
+ def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
+ hidden_state = self.proj(hidden_state)
+ hidden_state = self.act1(self.post_projection_norm(hidden_state))
+ return self.down_proj(self.act_fn(self.gate_proj(hidden_state)) * self.up_proj(hidden_state))
+
+
+class GlmOcrVisionPatchEmbed(nn.Module):
+ def __init__(self, config: GlmOcrVisionConfig) -> None:
+ super().__init__()
+ self.patch_size = config.patch_size
+ self.temporal_patch_size = config.temporal_patch_size
+ self.in_channels = config.in_channels
+ self.embed_dim = config.hidden_size
+
+ kernel_size = [self.temporal_patch_size, self.patch_size, self.patch_size]
+ self.proj = nn.Conv3d(self.in_channels, self.embed_dim, kernel_size=kernel_size, stride=kernel_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ target_dtype = self.proj.weight.dtype
+ hidden_states = hidden_states.view(
+ -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size
+ )
+ hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)
+ return hidden_states
+
+
+class GlmOcrVisionModel(GlmOcrPreTrainedModel):
+ config: GlmOcrVisionConfig
+ input_modalities = ("image", "video")
+ _no_split_modules = ["GlmOcrVisionBlock"]
+ _can_record_outputs = {
+ "hidden_states": GlmOcrVisionBlock,
+ "attentions": GlmOcrVisionAttention,
+ }
+
+ def __init__(self, config) -> None:
+ super().__init__(config)
+ self.spatial_merge_size = config.spatial_merge_size
+ self.patch_size = config.patch_size
+ self.patch_embed = GlmOcrVisionPatchEmbed(config)
+
+ head_dim = config.hidden_size // config.num_heads
+ self.rotary_pos_emb = GlmOcrVisionRotaryEmbedding(head_dim // 2)
+
+ self.blocks = nn.ModuleList([GlmOcrVisionBlock(config) for _ in range(config.depth)])
+ self.merger = GlmOcrVisionPatchMerger(
+ dim=config.out_hidden_size,
+ context_dim=config.out_hidden_size * config.in_channels,
+ hidden_act=config.hidden_act,
+ )
+ self.downsample = nn.Conv2d(
+ in_channels=config.hidden_size,
+ out_channels=config.out_hidden_size,
+ kernel_size=config.spatial_merge_size,
+ stride=config.spatial_merge_size,
+ )
+ self.post_layernorm = GlmOcrRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+ self.gradient_checkpointing = False
+ self.post_init()
+
+ def rot_pos_emb(self, grid_thw):
+ pos_ids = []
+ for t, h, w in grid_thw:
+ hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)
+ hpos_ids = hpos_ids.reshape(
+ h // self.spatial_merge_size,
+ self.spatial_merge_size,
+ w // self.spatial_merge_size,
+ self.spatial_merge_size,
+ )
+ hpos_ids = hpos_ids.permute(0, 2, 1, 3)
+ hpos_ids = hpos_ids.flatten()
+
+ wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)
+ wpos_ids = wpos_ids.reshape(
+ h // self.spatial_merge_size,
+ self.spatial_merge_size,
+ w // self.spatial_merge_size,
+ self.spatial_merge_size,
+ )
+ wpos_ids = wpos_ids.permute(0, 2, 1, 3)
+ wpos_ids = wpos_ids.flatten()
+ pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
+ pos_ids = torch.cat(pos_ids, dim=0)
+ max_grid_size = grid_thw[:, 1:].max()
+ rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
+ rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
+ return rotary_pos_emb, pos_ids
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor:
+ r"""
+ hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
+ The final hidden states of the model.
+ grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
+ The temporal, height and width of feature shape of each image in LLM.
+
+ Returns:
+ `torch.Tensor`: hidden_states.
+ """
+ hidden_states = self.patch_embed(hidden_states)
+ rotary_pos_emb, image_type_ids = self.rot_pos_emb(grid_thw)
+ emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+ position_embeddings = (emb.cos(), emb.sin())
+
+ cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+ dim=0,
+ # Select dtype based on the following factors:
+ # - FA2 requires that cu_seqlens_q must have dtype int32
+ # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw
+ # See https://github.com/huggingface/transformers/pull/34852 for more information
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+ )
+ cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+
+ for blk in self.blocks:
+ hidden_states = blk(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ position_embeddings=position_embeddings,
+ )
+
+ hidden_states = self.post_layernorm(hidden_states)
+
+ hidden_states = hidden_states.view(
+ -1, self.spatial_merge_size, self.spatial_merge_size, hidden_states.shape[-1]
+ )
+ hidden_states = hidden_states.permute(0, 3, 1, 2)
+ hidden_states = self.downsample(hidden_states).view(-1, self.config.out_hidden_size)
+
+ merged_hidden_states = self.merger(hidden_states)
+ return BaseModelOutputWithPooling(
+ last_hidden_state=hidden_states,
+ pooler_output=merged_hidden_states,
+ )
+
+
+class GlmOcrTextRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: GlmOcrTextConfig, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+ self.mrope_section = config.rope_parameters.get("mrope_section", [8, 12, 12])
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: GlmOcrTextConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
+ head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+ dim = int(head_dim * partial_rotary_factor)
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ # In contrast to other models, GLM-V has different position ids for the grids
+ # So we expand the inv_freq to shape (3, ...)
+ inv_freq_expanded = self.inv_freq[None, None, :, None].float().expand(3, position_ids.shape[1], -1, 1)
+ position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions)
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3)
+ freqs = self.apply_mrope(freqs, self.mrope_section)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+ def apply_mrope(self, freqs, mrope_section):
+ section = mrope_section
+ chunks = freqs.split(section, dim=-1)
+ result = torch.cat([chunk[i % 3] for i, chunk in enumerate(chunks)], dim=-1)
+ return result
+
+
+@auto_docstring
+class GlmOcrTextModel(GlmOcrPreTrainedModel):
+ config: GlmOcrTextConfig
+ input_modalities = ("text",)
+ _can_record_outputs = {
+ "hidden_states": GlmOcrTextDecoderLayer,
+ "attentions": GlmOcrTextAttention,
+ }
+
+ def __init__(self, config: GlmOcrTextConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList(
+ [GlmOcrTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = GlmOcrRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = GlmOcrTextRotaryEmbedding(config=config)
+
+ self.gradient_checkpointing = False
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple | BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ # torch.jit.trace() doesn't support cache objects in the output
+ if use_cache and past_key_values is None and not torch.jit.is_tracing():
+ past_key_values = DynamicCache(config=self.config)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ # the hard coded `3` is for temporal, height and width.
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1)
+ elif position_ids.ndim == 2:
+ position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1)
+
+ # NOTE: we need to pass text position ids for packing. Qwen2-VL uses 3D positions
+ # where each dim indicates visual spatial positions for temporal/height/width grids.
+ # There are two scenarios when FA2-like packed masking might be activated.
+ # 1. User specifically passed packed `position_ids` and no attention mask.
+ # In this case we expect the useer to create correct position ids for all 3 grids
+ # and prepend text-only position ids to it. The final tensor will be [4, bs, seq-len]
+ # 2. User runs forward with no attention mask and no position ids. In this case, position ids
+ # are prepared by the model (`get_rope_index`) as `[4, bs, seq-len]` tensor. Text-only positions are
+ # prepended by us when creating positions so that the mask is constructed correctly. NOTE: failing to pass
+ # text-only positions will cause incorrect mask construction, do not change `prepare_input_for_generation`
+ if position_ids.ndim == 3 and position_ids.shape[0] == 4:
+ text_position_ids = position_ids[0]
+ position_ids = position_ids[1:]
+ else:
+ # If inputs are not packed (usual 3D positions), do not prepare mask from position_ids
+ text_position_ids = None
+
+ mask_kwargs = {
+ "config": self.config,
+ "inputs_embeds": inputs_embeds,
+ "attention_mask": attention_mask,
+ "past_key_values": past_key_values,
+ "position_ids": text_position_ids,
+ }
+ # Create the masks
+ causal_mask = create_causal_mask(**mask_kwargs)
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for decoder_layer in self.layers:
+ layer_outputs = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_ids=text_position_ids,
+ past_key_values=past_key_values,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ hidden_states = layer_outputs
+
+ hidden_states = self.norm(hidden_states)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class GlmOcrModel(GlmOcrPreTrainedModel):
+ base_model_prefix = "model"
+ # Reference: fix gemma3 grad acc #37208
+ accepts_loss_kwargs = False
+ _no_split_modules = ["GlmOcrTextDecoderLayer", "GlmOcrVisionBlock"]
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.visual = GlmOcrVisionModel._from_config(config.vision_config)
+ self.language_model = GlmOcrTextModel._from_config(config.text_config)
+ self.rope_deltas = None # cache rope_deltas here
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.language_model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.language_model.set_input_embeddings(value)
+
+ def get_vision_position_ids(
+ self,
+ start_position: int,
+ grid_thw: list[int, int, int] | torch.Tensor,
+ temp_merge_size: int = 1,
+ spatial_merge_size: int = 1,
+ time_interval: int = 1,
+ device: str | torch.device | None = None,
+ ):
+ """
+ Compute 3D positional indices for vision tokens derived from a single image or video input.
+
+ The positions are generated from the input grid defined by temporal (T), height (H), and
+ width (W) dimensions. Temporal and spatial dimensions can be downscaled according to the
+ merge sizes used in the vision backbone. The resulting positions are offset by `start_position`.
+
+ Args:
+ start_position (`int`):
+ Offset added to all computed positional indices.
+ grid_thw (`Sequence[int]` or `torch.Tensor` of shape `(3,)`):
+ The (T, H, W) grid representing the feature layout of the current image or video after patch embedding.
+ temp_merge_size (`int`, *optional*):
+ Factor by which the temporal dimension is reduced in the backbone. The temporal grid size is divided
+ by this value. Defaults to 1.
+ spatial_merge_size (`int`, *optional*):
+ Factor by which the spatial dimensions (H and W) are reduced in the backbone. Both H and W are divided
+ by this value. Defaults to 1.
+ time_interval (`int`, *optional*):
+ Spacing factor applied between consecutive temporal position indices.Defaults to 1.
+ device (`str` or `torch.device`, *optional*):
+ Device on which the resulting tensor is allocated. If `None`, uses the current default device.
+
+ Returns:
+ torch.LongTensor of shape (3, sequence_length):
+ Positional indices for temporal, height, and width dimensions,
+ flattened into sequence form and offset by `start_position`.
+ """
+ llm_grid_t, llm_grid_h, llm_grid_w = (
+ grid_thw[0].item() // temp_merge_size,
+ grid_thw[1].item() // spatial_merge_size,
+ grid_thw[2].item() // spatial_merge_size,
+ )
+
+ image_seq_length = llm_grid_h * llm_grid_w * llm_grid_t
+ position_width = torch.arange(start_position, start_position + llm_grid_w, device=device).repeat(
+ llm_grid_h * llm_grid_t
+ )
+ position_height = torch.arange(start_position, start_position + llm_grid_h, device=device).repeat_interleave(
+ llm_grid_w * llm_grid_t
+ )
+ position_temporal = torch.full((image_seq_length,), start_position, device=device, dtype=torch.long)
+ position_temporal = position_temporal * time_interval
+ vision_position_ids = torch.stack([position_temporal, position_height, position_width], dim=0)
+
+ return vision_position_ids
+
+ def get_rope_index(
+ self,
+ input_ids: torch.LongTensor,
+ mm_token_type_ids: torch.IntTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Calculate the 3D rope index based on image and video's sizes. The utility expects a `vision + text`
+ sequence and will error out otherwise. For pure text sequence, please rely on model's auto-inferred
+ position ids. In a mixed vision + text sequence, vision tokens use 3D RoPE (temporal, height, width)
+ while text tokens use standard 1D RoPE.
+
+ Example:
+ Temporal patches: 3; Height patches: 2; Width patches: 2
+ Each vision input results in (temporal x height × width) positions. Here: 3 x 2 × 2 = 12 positions total.
+
+ Temporal position IDs are spaced by:
+ `interval = tokens_per_second * temporal_patch_size / fps`
+
+ If fps = 1; tokens_per_second = 25; temporal_patch_size = 2, temporal IDs increase by 50 for each temporal patch:
+ `[0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]`
+
+ Height IDs repeat per row: `[0, 0, 1, 1, ...]`
+ Width IDs alternate per column: `[0, 1, 0, 1, ...]`
+ Text tokens follow standard 1D RoPE and the position IDs grow consequently with a step of `1`
+
+ 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.
+ mm_token_type_ids (`torch.IntTensor` of shape `(batch_size, sequence_length)`):
+ Token type ids matching each modality to a different value in the input sequence, i.e. text (0), image (1), video (2).
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ 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**.
+
+ Returns:
+ position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)
+ mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)
+ """
+ spatial_merge_size = self.config.vision_config.spatial_merge_size
+
+ mrope_position_deltas = []
+ position_ids = torch.zeros(
+ 3,
+ input_ids.shape[0],
+ input_ids.shape[1],
+ dtype=input_ids.dtype,
+ device=input_ids.device,
+ )
+ grid_iters = {
+ 1: iter(image_grid_thw) if image_grid_thw is not None else None,
+ 2: iter(video_grid_thw) if video_grid_thw is not None else None,
+ }
+
+ for batch_idx, current_input_ids in enumerate(input_ids):
+ input_token_type = mm_token_type_ids[batch_idx]
+ if attention_mask is not None:
+ current_input_ids = current_input_ids[attention_mask[batch_idx].bool()]
+ input_token_type = input_token_type[attention_mask[batch_idx].bool()]
+
+ input_type_group = []
+ for key, group in itertools.groupby(enumerate(input_token_type.tolist()), lambda x: x[1]):
+ group = list(group)
+ start_index = group[0][0]
+ end_index = group[-1][0] + 1
+ input_type_group.append((key, start_index, end_index))
+
+ current_pos = 0
+ video_group_index = 0
+ llm_pos_ids_list = []
+ for modality_type, start_idx, end_idx in input_type_group:
+ # text == 0
+ if modality_type == 0:
+ text_len = end_idx - start_idx
+ llm_pos_ids_list.append(
+ torch.arange(text_len, device=input_ids.device).view(1, -1).expand(3, -1) + current_pos
+ )
+ current_pos += text_len
+ # image == 1, video == 2
+ else:
+ # GLM_OCR splits video into segments per frame but there's only one `grid_thw`
+ # per whole video. We can't exhaus the iterator and have to re-use the grid
+ # while processing the same video!
+ if modality_type == 2:
+ if video_group_index == 0:
+ grid_thw = next(grid_iters[modality_type])
+ video_group_index += 1
+ video_group_index = 0 if video_group_index >= grid_thw[0] else video_group_index
+ else:
+ grid_thw = next(grid_iters[modality_type])
+
+ # Videos are processed per frame separately, each temporal grid is always `1`
+ temp_merge_size = grid_thw[0]
+ vision_position_ids = self.get_vision_position_ids(
+ current_pos, grid_thw, temp_merge_size, spatial_merge_size, device=input_ids.device
+ )
+ llm_pos_ids_list.append(vision_position_ids)
+ current_pos += max(grid_thw[1], grid_thw[2]) // spatial_merge_size
+ llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+ if attention_mask is not None:
+ position_ids[:, batch_idx, attention_mask[batch_idx].bool()] = llm_positions.to(position_ids.device)
+ else:
+ position_ids[:, batch_idx] = llm_positions.to(position_ids.device)
+ mrope_position_deltas.append(llm_positions.max() + 1 - len(current_input_ids))
+ mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)
+ return position_ids, mrope_position_deltas
+
+ @can_return_tuple
+ @auto_docstring
+ def get_video_features(
+ self,
+ pixel_values_videos: torch.FloatTensor,
+ video_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input videos.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ """
+ pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
+ # reshape video_grid_thw -> [b, 3] -> [1, h, w] * frames
+ temp_frames_hw = []
+ video_grid_thw_list = video_grid_thw.tolist()
+ for t, h, w in video_grid_thw_list:
+ repeated_row = torch.tensor([1, h, w]).unsqueeze(0).repeat(t, 1)
+ temp_frames_hw.append(repeated_row)
+ flattened_video_grid_thw = torch.cat(temp_frames_hw, dim=0)
+ vision_outputs = self.visual(
+ pixel_values_videos, grid_thw=flattened_video_grid_thw, return_dict=True, **kwargs
+ )
+ split_sizes = (video_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
+ video_embeds = torch.split(vision_outputs.pooler_output, split_sizes)
+ vision_outputs.pooler_output = video_embeds
+
+ return vision_outputs
+
+ @can_return_tuple
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ pixel_values = pixel_values.type(self.visual.dtype)
+ vision_outputs = self.visual(pixel_values, grid_thw=image_grid_thw, **kwargs)
+ split_sizes = (image_grid_thw.prod(-1) // self.visual.spatial_merge_size**2).tolist()
+ image_embeds = torch.split(vision_outputs.pooler_output, split_sizes)
+ vision_outputs.pooler_output = image_embeds
+
+ return vision_outputs
+
+ def get_placeholder_mask(
+ self,
+ input_ids: torch.LongTensor,
+ inputs_embeds: torch.FloatTensor,
+ image_features: torch.FloatTensor | None = None,
+ video_features: torch.FloatTensor | None = None,
+ ):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`, and checks that the placeholder token count is
+ equal to the length of multimodal features. If the lengths are different, an error is raised.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ special_video_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.video_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_video_mask = special_video_mask.all(-1)
+ else:
+ # GLM-4.1V and GLM-4.5V special_video_mask is special_image_mask
+ special_image_mask = input_ids == self.config.image_token_id
+ special_video_mask = input_ids == self.config.image_token_id
+
+ n_image_tokens = special_image_mask.sum()
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if image_features is not None:
+ torch_compilable_check(
+ inputs_embeds[special_image_mask].numel() == image_features.numel(),
+ f"Image features and image tokens do not match, tokens: {n_image_tokens}, features: {image_features.shape[0]}",
+ )
+
+ n_video_tokens = special_video_mask.sum()
+ special_video_mask = special_video_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ if video_features is not None:
+ torch_compilable_check(
+ inputs_embeds[special_video_mask].numel() == video_features.numel(),
+ f"Video features and video tokens do not match, tokens: {n_video_tokens}, features: {video_features.shape[0]}",
+ )
+ return special_image_mask, special_video_mask
+
+ def compute_3d_position_ids(
+ self,
+ input_ids: torch.Tensor | None,
+ inputs_embeds: torch.Tensor | None,
+ image_grid_thw: torch.Tensor | None = None,
+ video_grid_thw: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: torch.Tensor | None = None,
+ mm_token_type_ids: torch.IntTensor | None = None,
+ ) -> torch.Tensor | None:
+ past_key_values_length = 0 if past_key_values is None else past_key_values.get_seq_length()
+ has_multimodal = image_grid_thw is not None or video_grid_thw is not None
+ if has_multimodal and mm_token_type_ids is None and input_ids is not None:
+ raise ValueError(
+ "Multimodal data was passed (via `image_grid_thw` or `video_grid_thw`) but `mm_token_type_ids` is "
+ "missing. Please pass `mm_token_type_ids` to the model so that multimodal RoPE (M-RoPE) can be "
+ "computed correctly. `mm_token_type_ids` is returned by the processor alongside `input_ids`."
+ )
+ can_compute_mrope = input_ids is not None and mm_token_type_ids is not None and has_multimodal
+
+ if can_compute_mrope and (self.rope_deltas is None or past_key_values_length == 0):
+ position_ids, rope_deltas = self.get_rope_index(
+ input_ids,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ attention_mask=attention_mask,
+ mm_token_type_ids=mm_token_type_ids,
+ )
+ self.rope_deltas = rope_deltas
+ # Use pre-calculated rope-deltas to infer correct 3D position ids during incremental
+ # generation (past_key_values_length > 0) or when only inputs_embeds is provided (no input_ids
+ # to recompute from). Skip when input_ids is provided without past_key_values to avoid shape
+ # mismatches from stale rope_deltas (e.g., training forward pass after generation).
+ elif self.rope_deltas is not None and (past_key_values_length > 0 or input_ids is None):
+ batch_size, seq_length, _ = inputs_embeds.shape
+ if attention_mask is not None:
+ position_ids = attention_mask.long().cumsum(-1) - 1
+ position_ids = position_ids.masked_fill(attention_mask == 0, 0)
+ position_ids = position_ids.view(1, batch_size, -1).repeat(3, 1, 1).to(inputs_embeds.device)
+ else:
+ position_ids = torch.arange(past_key_values_length, past_key_values_length + seq_length)
+ position_ids = position_ids.view(1, 1, -1).expand(3, batch_size, -1).to(inputs_embeds.device)
+ delta = self.rope_deltas.repeat_interleave(batch_size // self.rope_deltas.shape[0], dim=0)
+ position_ids = position_ids + delta.to(device=inputs_embeds.device)
+ else:
+ # Can't build correct 3D positions. Let the model infer it
+ position_ids = None
+ return position_ids
+
+ @auto_docstring
+ @can_return_tuple
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ pixel_values: torch.Tensor | None = None,
+ pixel_values_videos: torch.FloatTensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ rope_deltas: torch.LongTensor | None = None,
+ mm_token_type_ids: torch.IntTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | GlmOcrModelOutputWithPast:
+ r"""
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if pixel_values is not None:
+ image_embeds = self.get_image_features(pixel_values, image_grid_thw, return_dict=True).pooler_output
+ image_embeds = torch.cat(image_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ image_mask, _ = self.get_placeholder_mask(input_ids, inputs_embeds, image_features=image_embeds)
+ inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)
+
+ if pixel_values_videos is not None:
+ video_embeds = self.get_video_features(pixel_values_videos, video_grid_thw, return_dict=True).pooler_output
+ video_embeds = torch.cat(video_embeds, dim=0).to(inputs_embeds.device, inputs_embeds.dtype)
+ _, video_mask = self.get_placeholder_mask(input_ids, inputs_embeds, video_features=video_embeds)
+ inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)
+
+ if position_ids is None:
+ position_ids = self.compute_3d_position_ids(
+ input_ids=input_ids,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ mm_token_type_ids=mm_token_type_ids,
+ )
+
+ outputs = self.language_model(
+ input_ids=None,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ return GlmOcrModelOutputWithPast(
+ **outputs,
+ rope_deltas=self.rope_deltas,
+ )
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for GlmOcr causal language model (or autoregressive) outputs.
+ """
+)
+class GlmOcrCausalLMOutputWithPast(ModelOutput):
+ r"""
+ 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 (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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.
+ rope_deltas (`torch.LongTensor` of shape `(batch_size, )`, *optional*):
+ The rope index difference between sequence length and multimodal rope.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ rope_deltas: torch.LongTensor | None = None
+
+
+class GlmOcrForConditionalGeneration(GlmOcrPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}
+ # Reference: fix gemma3 grad acc #37208
+ accepts_loss_kwargs = False
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = GlmOcrModel(config)
+ self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False)
+
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.model.set_input_embeddings(value)
+
+ @auto_docstring
+ def get_video_features(
+ self,
+ pixel_values_videos: torch.FloatTensor,
+ video_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values_videos (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input videos.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+ """
+ return self.model.get_video_features(
+ pixel_values_videos=pixel_values_videos, video_grid_thw=video_grid_thw, **kwargs
+ )
+
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ image_grid_thw: torch.LongTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ """
+ return self.model.get_image_features(pixel_values=pixel_values, image_grid_thw=image_grid_thw, **kwargs)
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ pixel_values: torch.Tensor | None = None,
+ pixel_values_videos: torch.FloatTensor | None = None,
+ image_grid_thw: torch.LongTensor | None = None,
+ video_grid_thw: torch.LongTensor | None = None,
+ mm_token_type_ids: torch.IntTensor | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | GlmOcrCausalLMOutputWithPast:
+ r"""
+ 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]`.
+ image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):
+ The temporal, height and width of feature shape of each image in LLM.
+ video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):
+ The temporal, height and width of feature shape of each video in LLM.
+
+ Example:
+
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoProcessor, GlmOcrForConditionalGeneration
+
+ >>> model = GlmOcrForConditionalGeneration.from_pretrained("zai-org/GLM-4.1V-9B-Thinking")
+ >>> processor = AutoProcessor.from_pretrained("zai-org/GLM-4.1V-9B-Thinking")
+
+ >>> messages = [
+ {
+ "role": "user",
+ "content": [
+ {"type": "image", "url": "https://www.ilankelman.org/stopsigns/australia.jpg"},
+ {"type": "text", "text": "What is shown in this image?"},
+ ],
+ },
+ ]
+ >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+ >>> inputs = processor(text=[text], images=[image], vision_infos=[vision_infos])
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "The image shows a street scene with a red stop sign in the foreground. In the background, there is a large red gate with Chinese characters ..."
+ ```"""
+ outputs = self.model(
+ input_ids=input_ids,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ mm_token_type_ids=mm_token_type_ids,
+ position_ids=position_ids,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ hidden_states = outputs[0]
+
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size)
+
+ return GlmOcrCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ rope_deltas=outputs.rope_deltas,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ attention_mask=None,
+ inputs_embeds=None,
+ position_ids=None,
+ use_cache=True,
+ pixel_values=None,
+ pixel_values_videos=None,
+ image_grid_thw=None,
+ video_grid_thw=None,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ position_ids=position_ids,
+ pixel_values=pixel_values,
+ pixel_values_videos=pixel_values_videos,
+ image_grid_thw=image_grid_thw,
+ video_grid_thw=video_grid_thw,
+ use_cache=use_cache,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ if not is_first_iteration and use_cache:
+ model_inputs["pixel_values"] = None
+ model_inputs["pixel_values_videos"] = None
+
+ return model_inputs
+
+ def _prepare_position_ids_for_generation(self, inputs_tensor, model_kwargs):
+ # Overwritten -- requires 3D position ids
+
+ text_positions = super()._prepare_position_ids_for_generation(inputs_tensor, model_kwargs)
+
+ # Early exit in case we are continuing generation from past kv
+ past_length = 0
+ if (cache := model_kwargs.get("past_key_values")) is not None:
+ past_length = cache.get_seq_length()
+ if past_length != 0 and self.model.rope_deltas is not None:
+ position_ids = text_positions[None, ...] + self.model.rope_deltas
+ return position_ids
+
+ # Otherwise compute 3d position ids for vision tokens and concat with text position ids
+ if "input_ids" in model_kwargs and model_kwargs["input_ids"].shape[1] > 0:
+ inputs_tensor = model_kwargs["input_ids"]
+
+ is_input_ids = len(inputs_tensor.shape) == 2 and inputs_tensor.dtype in [torch.int, torch.long]
+ if (
+ is_input_ids
+ and model_kwargs.get("mm_token_type_ids") is not None
+ and (model_kwargs.get("image_grid_thw") is not None or model_kwargs.get("video_grid_thw") is not None)
+ ):
+ model_kwargs = {k: v for k, v in model_kwargs.items() if k != "input_ids"}
+ vision_positions, rope_deltas = self.model.get_rope_index(inputs_tensor, **model_kwargs)
+ self.model.rope_deltas = rope_deltas
+ else:
+ vision_positions = text_positions.unsqueeze(0).expand(3, -1, -1)
+ self.model.rope_deltas = torch.zeros(
+ inputs_tensor.shape[0], 1, dtype=torch.long, device=inputs_tensor.device
+ )
+
+ # Concatenate "text + vision" positions into [4, bs, seq-len]
+ text_positions = text_positions[None, ...]
+ position_ids = torch.cat([text_positions, vision_positions], dim=0)
+
+ return position_ids
+
+ def _get_image_nums_and_video_nums(
+ self,
+ input_ids: torch.LongTensor | None,
+ inputs_embeds: torch.Tensor | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ """
+ Get the number of images and videos for each sample to calculate the separation length of the sample tensor.
+ These parameters are not passed through the processor to avoid unpredictable impacts from interface modifications.
+
+ Args:
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
+ Indices of input sequence tokens in the vocabulary.
+
+ Returns:
+ image_nums (`torch.LongTensor` of shape `(batch_size, num_images_sample)`)
+ video_nums (`torch.LongTensor` of shape `(batch_size, num_videos_sample)`)
+ """
+
+ if inputs_embeds is not None:
+ is_image = (
+ inputs_embeds
+ == self.get_input_embeddings()(
+ torch.tensor(self.config.image_start_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ )[..., 0]
+ is_video_start = (
+ inputs_embeds
+ == self.get_input_embeddings()(
+ torch.tensor(self.config.video_start_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ )[..., 0]
+ is_video_end = (
+ inputs_embeds
+ == self.get_input_embeddings()(
+ torch.tensor(self.config.video_end_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ )[..., 0]
+ else:
+ is_image = input_ids == self.config.image_start_token_id
+ is_video_start = input_ids == self.config.video_start_token_id
+ is_video_end = input_ids == self.config.video_end_token_id
+
+ # Cumulative sum to track if we're inside a video span
+ # We'll assume well-formed video tags (i.e. matching starts and ends)
+ video_level = torch.cumsum(is_video_start.int() - is_video_end.int(), dim=1)
+ inside_video = video_level > 0 # shape (batch_size, seq_length)
+
+ # Mask out image tokens that are inside video spans
+ standalone_images = is_image & (~inside_video)
+
+ # Count per batch
+ image_counts = standalone_images.sum(dim=1)
+ video_counts = is_video_start.sum(dim=1)
+
+ return image_counts, video_counts
+
+ def _expand_inputs_for_generation(
+ self,
+ expand_size: int = 1,
+ is_encoder_decoder: bool = False,
+ input_ids: torch.LongTensor | None = None,
+ **model_kwargs,
+ ) -> tuple[torch.LongTensor, dict[str, Any]]:
+ # Overwritten -- Support for expanding tensors without a batch size dimension
+ # e.g., pixel_values, image_grid_thw, pixel_values_videos, video_grid_thw, second_per_grid_t
+ # pixel_values.shape[0] is sum(seqlen_images for samples)
+ # image_grid_thw.shape[0] is sum(num_images for samples)
+
+ if expand_size == 1:
+ return input_ids, model_kwargs
+
+ visual_keys = ["pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw", "second_per_grid_ts"]
+
+ def _expand_dict_for_generation_visual(dict_to_expand):
+ image_grid_thw = model_kwargs.get("image_grid_thw", None)
+ video_grid_thw = model_kwargs.get("video_grid_thw", None)
+ image_nums, video_nums = self._get_image_nums_and_video_nums(
+ input_ids, inputs_embeds=model_kwargs.get("inputs_embeds", None)
+ )
+
+ def _repeat_interleave_samples(x, lengths, repeat_times):
+ samples = torch.split(x, lengths)
+ repeat_args = [repeat_times] + [1] * (x.dim() - 1)
+ result = torch.cat([sample.repeat(*repeat_args) for sample in samples], dim=0)
+ return result
+
+ for key in dict_to_expand:
+ if key == "pixel_values":
+ # split images into samples
+ samples = torch.split(image_grid_thw, list(image_nums))
+ # compute the sequence length of images for each sample
+ lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "image_grid_thw":
+ # get the num of images for each sample
+ lengths = list(image_nums)
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "pixel_values_videos":
+ samples = torch.split(video_grid_thw, list(video_nums))
+ lengths = [torch.prod(sample, dim=1).sum() for sample in samples]
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "video_grid_thw":
+ lengths = list(video_nums)
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=lengths, repeat_times=expand_size
+ )
+ elif key == "second_per_grid_ts":
+ dict_to_expand[key] = _repeat_interleave_samples(
+ dict_to_expand[key], lengths=list(video_nums), repeat_times=expand_size
+ )
+ return dict_to_expand
+
+ def _expand_dict_for_generation(dict_to_expand):
+ for key in dict_to_expand:
+ if key == "position_ids" and dict_to_expand[key].ndim == 3:
+ dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=1)
+ elif (
+ dict_to_expand[key] is not None
+ and isinstance(dict_to_expand[key], torch.Tensor)
+ and key not in visual_keys
+ ):
+ dict_to_expand[key] = dict_to_expand[key].repeat_interleave(expand_size, dim=0)
+ return dict_to_expand
+
+ model_kwargs = _expand_dict_for_generation_visual(model_kwargs)
+
+ if input_ids is not None:
+ input_ids = input_ids.repeat_interleave(expand_size, dim=0)
+
+ model_kwargs = _expand_dict_for_generation(model_kwargs)
+
+ if is_encoder_decoder:
+ if model_kwargs.get("encoder_outputs") is None:
+ raise ValueError("If `is_encoder_decoder` is True, make sure that `encoder_outputs` is defined.")
+ model_kwargs["encoder_outputs"] = _expand_dict_for_generation(model_kwargs["encoder_outputs"])
+
+ return input_ids, model_kwargs
+
+
+__all__ = [
+ "GlmOcrTextModel",
+ "GlmOcrVisionModel",
+ "GlmOcrModel",
+ "GlmOcrPreTrainedModel",
+ "GlmOcrForConditionalGeneration",
+]
diff --git a/third_party/transformers/src/transformers/models/glm_ocr/modular_glm_ocr.py b/third_party/transformers/src/transformers/models/glm_ocr/modular_glm_ocr.py
new file mode 100644
index 0000000000000000000000000000000000000000..2f71dded711d8fbcfb5ab88f83944be96b168910
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/glm_ocr/modular_glm_ocr.py
@@ -0,0 +1,321 @@
+# Copyright 2026 the HuggingFace Team. All rights reserved.
+#
+# 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.
+
+from collections.abc import Callable
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from huggingface_hub.dataclasses import strict
+
+from ...modeling_outputs import BaseModelOutputWithPooling
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
+from ...utils import auto_docstring
+from ..glm4v.configuration_glm4v import Glm4vConfig, Glm4vTextConfig, Glm4vVisionConfig
+from ..glm4v.modeling_glm4v import (
+ Glm4vForConditionalGeneration,
+ Glm4VisionMlp,
+ Glm4vModel,
+ Glm4vModelOutputWithPast,
+ Glm4vPreTrainedModel,
+ Glm4vRMSNorm,
+ Glm4vTextAttention,
+ Glm4vTextDecoderLayer,
+ Glm4vTextModel,
+ Glm4vVisionAttention,
+ Glm4vVisionBlock,
+ Glm4vVisionModel,
+ Glm4vVisionPatchMerger,
+ apply_rotary_pos_emb_vision,
+ eager_attention_forward,
+ is_flash_attention_requested,
+)
+
+
+class GlmOcrRMSNorm(Glm4vRMSNorm):
+ pass
+
+
+class GlmOcrVisionMlp(Glm4VisionMlp):
+ def __init__(self, config, bias: bool = True):
+ super().__init__(config)
+ self.intermediate_size = config.intermediate_size
+
+
+@auto_docstring(checkpoint="zai-org/GLM-OCR")
+@strict
+class GlmOcrVisionConfig(Glm4vVisionConfig):
+ hidden_size: int = 1024
+ attention_bias: bool = True
+ num_heads: int = 16
+ out_hidden_size: int = 1536
+ intermediate_size: int = 4096
+
+
+@auto_docstring(checkpoint="zai-org/GLM-OCR")
+@strict
+class GlmOcrTextConfig(Glm4vTextConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import GlmOcrTextModel, GlmOcrConfig
+
+ >>> # Initializing a GLM-OCR style configuration
+ >>> configuration = GlmOcrConfig()
+
+ >>> # Initializing a model from the GLM-OCR style configuration
+ >>> model = GlmOcrTextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ vocab_size: int = 59392
+ hidden_size: int = 1024
+ intermediate_size: int = 4096
+ num_hidden_layers: int = 16
+ num_attention_heads: int = 16
+ num_key_value_heads: int = 8
+ max_position_embeddings: int = 131072
+
+
+@auto_docstring(checkpoint="zai-org/GLM-OCR")
+@strict
+class GlmOcrConfig(Glm4vConfig):
+ r"""
+ image_start_token_id (`int`, *optional*, defaults to 59256):
+ The image start token index to encode the start of image.
+ image_end_token_id (`int`, *optional*, defaults to 59257):
+ The image end token index to encode the end of image.
+ video_start_token_id (`int`, *optional*, defaults to 59258):
+ The video start token index to encode the start of video.
+ video_end_token_id (`int`, *optional*, defaults to 59259):
+ The video end token index to encode the end of video.
+
+ ```python
+ >>> from transformers import GlmOcrForConditionalGeneration, GlmOcrConfig
+
+ >>> # Initializing a GLM-OCR style configuration
+ >>> configuration = GlmOcrConfig()
+
+ >>> # Initializing a model from the GLM-OCR style configuration
+ >>> model = GlmOcrForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ image_token_id: int = 59280
+ video_token_id: int = 59281
+ image_start_token_id: int = 59256
+ image_end_token_id: int = 59257
+ video_start_token_id: int = 59258
+ video_end_token_id: int = 59259
+
+
+class GlmOcrTextAttention(Glm4vTextAttention, nn.Module):
+ def __init__(self, config: GlmOcrTextConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_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)
+
+
+class GlmOcrTextDecoderLayer(Glm4vTextDecoderLayer):
+ pass
+
+
+class GlmOcrPreTrainedModel(Glm4vPreTrainedModel):
+ _keys_to_ignore_on_load_unexpected = [r"model\.language_model\.layers\.16.*"]
+
+
+class GlmOcrModelOutputWithPast(Glm4vModelOutputWithPast):
+ pass
+
+
+class GlmOcrVisionAttention(Glm4vVisionAttention):
+ def __init__(self, config: GlmOcrVisionConfig) -> None:
+ super().__init__()
+ self.qkv = nn.Linear(config.hidden_size, config.hidden_size * 3, bias=config.attention_bias)
+ self.proj = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)
+ self.q_norm = GlmOcrRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+ self.k_norm = GlmOcrRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ rotary_pos_emb: torch.Tensor | None = None,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ query_states, key_states, value_states = (
+ self.qkv(hidden_states).reshape(seq_length, 3, self.num_heads, -1).permute(1, 0, 2, 3).unbind(0)
+ )
+
+ query_states = self.q_norm(query_states)
+ key_states = self.k_norm(key_states)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb_vision(query_states, key_states, cos, sin)
+ query_states = query_states.transpose(0, 1).unsqueeze(0)
+ key_states = key_states.transpose(0, 1).unsqueeze(0)
+ value_states = value_states.transpose(0, 1).unsqueeze(0)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ if is_flash_attention_requested(self.config):
+ # Flash Attention: Use cu_seqlens for variable length attention
+ max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()
+ attn_output, _ = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ cu_seq_lens_q=cu_seqlens,
+ cu_seq_lens_k=cu_seqlens,
+ max_length_q=max_seqlen,
+ max_length_k=max_seqlen,
+ is_causal=False,
+ **kwargs,
+ )
+ else:
+ # Other implementations: Process each chunk separately
+ lengths = cu_seqlens[1:] - cu_seqlens[:-1]
+ splits = [
+ torch.split(tensor, lengths.tolist(), dim=2) for tensor in (query_states, key_states, value_states)
+ ]
+
+ attn_outputs = [
+ attention_interface(
+ self,
+ q,
+ k,
+ v,
+ attention_mask=None,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ is_causal=False,
+ **kwargs,
+ )[0]
+ for q, k, v in zip(*splits)
+ ]
+ attn_output = torch.cat(attn_outputs, dim=1)
+
+ attn_output = attn_output.reshape(seq_length, -1).contiguous()
+ attn_output = self.proj(attn_output)
+ return attn_output
+
+
+class GlmOcrVisionBlock(Glm4vVisionBlock):
+ def __init__(self, config) -> None:
+ super().__init__()
+ self.mlp = GlmOcrVisionMlp(config, bias=config.attention_bias)
+
+
+class GlmOcrVisionPatchMerger(Glm4vVisionPatchMerger):
+ pass
+
+
+class GlmOcrVisionModel(Glm4vVisionModel):
+ def __init__(self, config) -> None:
+ super().__init__(config)
+ del self.embeddings
+ del self.post_conv_layernorm
+ self.merger = GlmOcrVisionPatchMerger(
+ dim=config.out_hidden_size,
+ context_dim=config.out_hidden_size * config.in_channels,
+ hidden_act=config.hidden_act,
+ )
+
+ def forward(self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, **kwargs) -> torch.Tensor:
+ r"""
+ hidden_states (`torch.Tensor` of shape `(seq_len, hidden_size)`):
+ The final hidden states of the model.
+ grid_thw (`torch.Tensor` of shape `(num_images_or_videos, 3)`):
+ The temporal, height and width of feature shape of each image in LLM.
+
+ Returns:
+ `torch.Tensor`: hidden_states.
+ """
+ hidden_states = self.patch_embed(hidden_states)
+ rotary_pos_emb, image_type_ids = self.rot_pos_emb(grid_thw)
+ emb = torch.cat((rotary_pos_emb, rotary_pos_emb), dim=-1)
+ position_embeddings = (emb.cos(), emb.sin())
+
+ cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(
+ dim=0,
+ # Select dtype based on the following factors:
+ # - FA2 requires that cu_seqlens_q must have dtype int32
+ # - torch.onnx.export requires that cu_seqlens_q must have same dtype as grid_thw
+ # See https://github.com/huggingface/transformers/pull/34852 for more information
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
+ )
+ cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
+
+ for blk in self.blocks:
+ hidden_states = blk(
+ hidden_states,
+ cu_seqlens=cu_seqlens,
+ position_embeddings=position_embeddings,
+ )
+
+ hidden_states = self.post_layernorm(hidden_states)
+
+ hidden_states = hidden_states.view(
+ -1, self.spatial_merge_size, self.spatial_merge_size, hidden_states.shape[-1]
+ )
+ hidden_states = hidden_states.permute(0, 3, 1, 2)
+ hidden_states = self.downsample(hidden_states).view(-1, self.config.out_hidden_size)
+
+ merged_hidden_states = self.merger(hidden_states)
+ return BaseModelOutputWithPooling(
+ last_hidden_state=hidden_states,
+ pooler_output=merged_hidden_states,
+ )
+
+
+class GlmOcrTextModel(Glm4vTextModel):
+ _can_record_outputs = {
+ "hidden_states": GlmOcrTextDecoderLayer,
+ "attentions": GlmOcrTextAttention,
+ }
+
+
+class GlmOcrModel(Glm4vModel):
+ pass
+
+
+class GlmOcrForConditionalGeneration(Glm4vForConditionalGeneration):
+ pass
+
+
+__all__ = [
+ "GlmOcrConfig",
+ "GlmOcrTextConfig",
+ "GlmOcrVisionConfig",
+ "GlmOcrTextModel",
+ "GlmOcrVisionModel",
+ "GlmOcrModel",
+ "GlmOcrPreTrainedModel",
+ "GlmOcrForConditionalGeneration",
+]
diff --git a/third_party/transformers/src/transformers/models/gpt_neox/__init__.py b/third_party/transformers/src/transformers/models/gpt_neox/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..91ef201d1126c45166aa8256cdc0e07de0444e56
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/gpt_neox/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_gpt_neox import *
+ from .modeling_gpt_neox import *
+ from .tokenization_gpt_neox import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/gpt_neox/configuration_gpt_neox.py b/third_party/transformers/src/transformers/models/gpt_neox/configuration_gpt_neox.py
new file mode 100644
index 0000000000000000000000000000000000000000..256ddca26b4921ad8680a40caee2455a028afddc
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/gpt_neox/configuration_gpt_neox.py
@@ -0,0 +1,103 @@
+# Copyright 2022 EleutherAI and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""GPTNeoX model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="EleutherAI/gpt-neox-20b")
+@strict
+class GPTNeoXConfig(PreTrainedConfig):
+ r"""
+ use_parallel_residual (`bool`, *optional*, defaults to `True`):
+ Whether to use a "parallel" formulation in each Transformer layer, which can provide a slight training
+ speedup at large scales (e.g. 20B).
+
+ Example:
+
+ ```python
+ >>> from transformers import GPTNeoXConfig, GPTNeoXModel
+
+ >>> # Initializing a GPTNeoX gpt-neox-20b style configuration
+ >>> configuration = GPTNeoXConfig()
+
+ >>> # Initializing a model (with random weights) from the gpt-neox-20b style configuration
+ >>> model = GPTNeoXModel(configuration) # doctest: +SKIP
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config # doctest: +SKIP
+ ```"""
+
+ model_type = "gpt_neox"
+ keys_to_ignore_at_inference = ["past_key_values"]
+ base_model_tp_plan = {
+ "layers.*.attention.query_key_value": "colwise",
+ "layers.*.attention.dense": "rowwise",
+ "layers.*.mlp.dense_h_to_4h": "colwise",
+ "layers.*.mlp.dense_4h_to_h": "rowwise",
+ }
+ base_model_pp_plan = {
+ "embed_in": (["input_ids"], ["inputs_embeds"]),
+ "emb_dropout": (["inputs_embeds"], ["hidden_states"]),
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
+ "final_layer_norm": (["hidden_states"], ["hidden_states"]),
+ }
+
+ vocab_size: int = 50432
+ hidden_size: int = 6144
+ num_hidden_layers: int = 44
+ num_attention_heads: int = 64
+ intermediate_size: int = 24576
+ hidden_act: str = "gelu"
+ attention_dropout: float | int = 0.0
+ hidden_dropout: float | int = 0.0
+ classifier_dropout: float | int = 0.1
+ max_position_embeddings: int = 2048
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-5
+ use_cache: bool = True
+ bos_token_id: int | None = 0
+ eos_token_id: int | list[int] | None = 2
+ pad_token_id: int | None = None
+ tie_word_embeddings: bool = False
+ use_parallel_residual: bool = True
+ rope_parameters: RopeParameters | dict | None = None
+ attention_bias: bool = True
+ is_decoder: bool = False
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ if self.hidden_size % self.num_attention_heads != 0:
+ raise ValueError(
+ "The hidden size is not divisible by the number of attention heads! Make sure to update them!"
+ )
+
+ def convert_rope_params_to_dict(self, **kwargs):
+ rope_scaling = kwargs.pop("rope_scaling", None)
+ self.rope_parameters = rope_scaling or self.rope_parameters
+ self.rope_parameters = self.rope_parameters if self.rope_parameters is not None else {}
+
+ # Standardize and validate the correctness of rotary position embeddings parameters
+ # Model uses non-standard naming for rope params, overwrite!
+ self.rope_parameters.setdefault("rope_theta", kwargs.pop("rotary_emb_base", self.default_theta))
+ self.rope_parameters.setdefault("partial_rotary_factor", kwargs.pop("rotary_pct", 0.25))
+ self.standardize_rope_params()
+ return kwargs
+
+
+__all__ = ["GPTNeoXConfig"]
diff --git a/third_party/transformers/src/transformers/models/gpt_neox/modeling_gpt_neox.py b/third_party/transformers/src/transformers/models/gpt_neox/modeling_gpt_neox.py
new file mode 100644
index 0000000000000000000000000000000000000000..10e4b5922add0d9a2a589f82b488c27b05c497e8
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/gpt_neox/modeling_gpt_neox.py
@@ -0,0 +1,686 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/gpt_neox/modular_gpt_neox.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_gpt_neox.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutputWithPast,
+ CausalLMOutputWithPast,
+ QuestionAnsweringModelOutput,
+ SequenceClassifierOutputWithPast,
+ TokenClassifierOutput,
+)
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils.generic import maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_gpt_neox import GPTNeoXConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class GPTNeoXMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense_h_to_4h = nn.Linear(config.hidden_size, config.intermediate_size)
+ self.dense_4h_to_h = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.act = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_states):
+ hidden_states = self.dense_h_to_4h(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.dense_4h_to_h(hidden_states)
+ return hidden_states
+
+
+class GPTNeoXRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: GPTNeoXConfig, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: GPTNeoXConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
+ head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+ dim = int(head_dim * partial_rotary_factor)
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+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)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+
+ # Keep half or full tensor for later concatenation
+ rotary_dim = cos.shape[-1]
+ q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+ k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+ # Apply rotary embeddings on the first half or full tensor
+ q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+ k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+
+ # Concatenate back to full shape
+ q_embed = torch.cat([q_embed, q_pass], dim=-1)
+ k_embed = torch.cat([k_embed, k_pass], dim=-1)
+ return q_embed, k_embed
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor,
+ scaling: float,
+ dropout: float | int = 0.0,
+ **kwargs,
+):
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value)
+
+ # Reshape outputs
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class GPTNeoXAttention(nn.Module):
+ def __init__(self, config, layer_idx=None):
+ super().__init__()
+ self.config = config
+ self.head_size = config.hidden_size // config.num_attention_heads
+ self.attention_dropout = config.attention_dropout
+ partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
+ self.rotary_ndims = int(self.head_size * partial_rotary_factor)
+ self.scaling = self.head_size**-0.5
+ self.is_causal = True
+ self.layer_idx = layer_idx
+
+ self.query_key_value = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=config.attention_bias)
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor,
+ attention_mask: torch.FloatTensor,
+ layer_past: Cache | None = None,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ):
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, 3 * self.head_size)
+
+ qkv = self.query_key_value(hidden_states).view(hidden_shape).transpose(1, 2)
+ query_states, key_states, value_states = qkv.chunk(3, dim=-1)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ # Cache QKV values
+ if layer_past is not None:
+ key_states, value_states = layer_past.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ # Compute attention
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ **kwargs,
+ )
+
+ # Reshape outputs and final projection
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.dense(attn_output)
+
+ return attn_output, attn_weights
+
+
+class GPTNeoXLayer(GradientCheckpointingLayer):
+ def __init__(self, config, layer_idx):
+ super().__init__()
+ self.use_parallel_residual = config.use_parallel_residual
+ self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.post_attention_dropout = nn.Dropout(config.hidden_dropout)
+ self.post_mlp_dropout = nn.Dropout(config.hidden_dropout)
+ self.attention = GPTNeoXAttention(config, layer_idx)
+ self.mlp = GPTNeoXMLP(config)
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor | None,
+ attention_mask: torch.FloatTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ use_cache: bool | None = False,
+ layer_past: Cache | None = None,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ):
+ attn_output, _ = self.attention(
+ self.input_layernorm(hidden_states),
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ layer_past=layer_past,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ attn_output = self.post_attention_dropout(attn_output)
+
+ if self.use_parallel_residual:
+ mlp_output = self.mlp(self.post_attention_layernorm(hidden_states))
+ mlp_output = self.post_mlp_dropout(mlp_output)
+ hidden_states = mlp_output + attn_output + hidden_states
+ else:
+ attn_output = attn_output + hidden_states
+ mlp_output = self.mlp(self.post_attention_layernorm(attn_output))
+ mlp_output = self.post_mlp_dropout(mlp_output)
+ hidden_states = mlp_output + attn_output
+
+ return hidden_states
+
+
+@auto_docstring
+class GPTNeoXPreTrainedModel(PreTrainedModel):
+ config: GPTNeoXConfig
+ base_model_prefix = "gpt_neox"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["GPTNeoXLayer"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": GPTNeoXLayer,
+ "attentions": GPTNeoXAttention,
+ }
+ _keys_to_ignore_on_load_unexpected = [r"attention.bias", r"attention.masked_bias"]
+
+
+@auto_docstring
+class GPTNeoXModel(GPTNeoXPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.config = config
+
+ self.embed_in = nn.Embedding(config.vocab_size, config.hidden_size)
+ self.emb_dropout = nn.Dropout(config.hidden_dropout)
+ self.layers = nn.ModuleList([GPTNeoXLayer(config, i) for i in range(config.num_hidden_layers)])
+ self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.rotary_emb = GPTNeoXRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_in(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ hidden_states = self.emb_dropout(inputs_embeds)
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+ for layer in self.layers:
+ hidden_states = layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_ids=position_ids,
+ layer_past=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = self.final_layer_norm(hidden_states)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+ def get_input_embeddings(self):
+ return self.embed_in
+
+ def set_input_embeddings(self, value):
+ self.embed_in = value
+
+
+@auto_docstring(
+ custom_intro="""
+ GPTNeoX Model with a `language modeling` head on top for CLM fine-tuning.
+ """
+)
+class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"embed_out.weight": "gpt_neox.embed_in.weight"}
+ _tp_plan = {"embed_out": "colwise_gather_output"}
+ _pp_plan = {"embed_out": (["hidden_states"], ["logits"])}
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.gpt_neox = GPTNeoXModel(config)
+ self.embed_out = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.embed_out
+
+ def set_output_embeddings(self, new_embeddings):
+ self.embed_out = new_embeddings
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithPast:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
+ `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
+ ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, GPTNeoXForCausalLM, GPTNeoXConfig
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
+ >>> config = GPTNeoXConfig.from_pretrained("EleutherAI/gpt-neox-20b")
+ >>> config.is_decoder = True
+ >>> model = GPTNeoXForCausalLM.from_pretrained("EleutherAI/gpt-neox-20b", config=config)
+
+ >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+ >>> outputs = model(**inputs)
+
+ >>> prediction_logits = outputs.logits
+ ```"""
+
+ outputs: BaseModelOutputWithPast = self.gpt_neox(
+ input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.embed_out(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The GPTNeoX Model transformer with a sequence classification head on top (linear layer).
+
+ [`GPTNeoXForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+ (e.g. GPT-1) do.
+
+ Since it does classification on the last token, it requires to know the position of the last token. If a
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+ each row of the batch).
+ """
+)
+class GPTNeoXForSequenceClassification(GPTNeoXPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.gpt_neox = GPTNeoXModel(config)
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs,
+ ) -> SequenceClassifierOutputWithPast:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+
+ outputs: BaseModelOutputWithPast = self.gpt_neox(
+ input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+ hidden_states = outputs.last_hidden_state
+ logits = self.score(hidden_states)
+
+ batch_size = logits.shape[0]
+ if self.config.pad_token_id is None and batch_size != 1:
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+ if self.config.pad_token_id is None:
+ last_non_pad_token = -1
+ elif input_ids is not None:
+ # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
+ non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+ token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+ last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+ else:
+ last_non_pad_token = -1
+ logger.warning_once(
+ f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+ "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+ )
+
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)
+
+ return SequenceClassifierOutputWithPast(
+ loss=loss,
+ logits=pooled_logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+class GPTNeoXForTokenClassification(GPTNeoXPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.gpt_neox = GPTNeoXModel(config)
+ self.dropout = nn.Dropout(config.classifier_dropout)
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs,
+ ) -> TokenClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+
+ outputs: BaseModelOutputWithPast = self.gpt_neox(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ hidden_states = self.dropout(hidden_states)
+ logits = self.classifier(hidden_states)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits, labels, self.config)
+
+ return TokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring
+class GPTNeoXForQuestionAnswering(GPTNeoXPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.gpt_neox = GPTNeoXModel(config)
+ self.qa_outputs = nn.Linear(config.hidden_size, 2)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ start_positions: torch.LongTensor | None = None,
+ end_positions: torch.LongTensor | None = None,
+ **kwargs,
+ ) -> QuestionAnsweringModelOutput:
+ outputs: BaseModelOutputWithPast = self.gpt_neox(
+ input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ sequence_output = outputs.last_hidden_state
+
+ logits = self.qa_outputs(sequence_output)
+ start_logits, end_logits = logits.split(1, dim=-1)
+ start_logits = start_logits.squeeze(-1).contiguous()
+ end_logits = end_logits.squeeze(-1).contiguous()
+
+ loss = None
+ if start_positions is not None and end_positions is not None:
+ loss = self.loss_function(start_logits, end_logits, start_positions, end_positions)
+
+ return QuestionAnsweringModelOutput(
+ loss=loss,
+ start_logits=start_logits,
+ end_logits=end_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "GPTNeoXForCausalLM",
+ "GPTNeoXForQuestionAnswering",
+ "GPTNeoXForSequenceClassification",
+ "GPTNeoXForTokenClassification",
+ "GPTNeoXLayer",
+ "GPTNeoXModel",
+ "GPTNeoXPreTrainedModel",
+]
diff --git a/third_party/transformers/src/transformers/models/gpt_neox/modular_gpt_neox.py b/third_party/transformers/src/transformers/models/gpt_neox/modular_gpt_neox.py
new file mode 100644
index 0000000000000000000000000000000000000000..f778501b7b38e3286871726bf3f0d7f6d647636c
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/gpt_neox/modular_gpt_neox.py
@@ -0,0 +1,633 @@
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_causal_mask
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutputWithPast,
+ CausalLMOutputWithPast,
+ QuestionAnsweringModelOutput,
+ SequenceClassifierOutputWithPast,
+ TokenClassifierOutput,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from ..llama.modeling_llama import LlamaModel, LlamaPreTrainedModel, LlamaRotaryEmbedding, rotate_half
+from .configuration_gpt_neox import GPTNeoXConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class GPTNeoXMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense_h_to_4h = nn.Linear(config.hidden_size, config.intermediate_size)
+ self.dense_4h_to_h = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.act = ACT2FN[config.hidden_act]
+
+ def forward(self, hidden_states):
+ hidden_states = self.dense_h_to_4h(hidden_states)
+ hidden_states = self.act(hidden_states)
+ hidden_states = self.dense_4h_to_h(hidden_states)
+ return hidden_states
+
+
+class GPTNeoXRotaryEmbedding(LlamaRotaryEmbedding):
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: GPTNeoXConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
+ head_dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+ dim = int(head_dim * partial_rotary_factor)
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+
+ # Keep half or full tensor for later concatenation
+ rotary_dim = cos.shape[-1]
+ q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+ k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+
+ # Apply rotary embeddings on the first half or full tensor
+ q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+ k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+
+ # Concatenate back to full shape
+ q_embed = torch.cat([q_embed, q_pass], dim=-1)
+ k_embed = torch.cat([k_embed, k_pass], dim=-1)
+ return q_embed, k_embed
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor,
+ scaling: float,
+ dropout: float | int = 0.0,
+ **kwargs,
+):
+ attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value)
+
+ # Reshape outputs
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class GPTNeoXAttention(nn.Module):
+ def __init__(self, config, layer_idx=None):
+ super().__init__()
+ self.config = config
+ self.head_size = config.hidden_size // config.num_attention_heads
+ self.attention_dropout = config.attention_dropout
+ partial_rotary_factor = config.rope_parameters.get("partial_rotary_factor", 1.0)
+ self.rotary_ndims = int(self.head_size * partial_rotary_factor)
+ self.scaling = self.head_size**-0.5
+ self.is_causal = True
+ self.layer_idx = layer_idx
+
+ self.query_key_value = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=config.attention_bias)
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor,
+ attention_mask: torch.FloatTensor,
+ layer_past: Cache | None = None,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ):
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, 3 * self.head_size)
+
+ qkv = self.query_key_value(hidden_states).view(hidden_shape).transpose(1, 2)
+ query_states, key_states, value_states = qkv.chunk(3, dim=-1)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ # Cache QKV values
+ if layer_past is not None:
+ key_states, value_states = layer_past.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ # Compute attention
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ scaling=self.scaling,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ **kwargs,
+ )
+
+ # Reshape outputs and final projection
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.dense(attn_output)
+
+ return attn_output, attn_weights
+
+
+class GPTNeoXLayer(GradientCheckpointingLayer):
+ def __init__(self, config, layer_idx):
+ super().__init__()
+ self.use_parallel_residual = config.use_parallel_residual
+ self.input_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.post_attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.post_attention_dropout = nn.Dropout(config.hidden_dropout)
+ self.post_mlp_dropout = nn.Dropout(config.hidden_dropout)
+ self.attention = GPTNeoXAttention(config, layer_idx)
+ self.mlp = GPTNeoXMLP(config)
+
+ def forward(
+ self,
+ hidden_states: torch.FloatTensor | None,
+ attention_mask: torch.FloatTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ use_cache: bool | None = False,
+ layer_past: Cache | None = None,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ):
+ attn_output, _ = self.attention(
+ self.input_layernorm(hidden_states),
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ layer_past=layer_past,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ attn_output = self.post_attention_dropout(attn_output)
+
+ if self.use_parallel_residual:
+ mlp_output = self.mlp(self.post_attention_layernorm(hidden_states))
+ mlp_output = self.post_mlp_dropout(mlp_output)
+ hidden_states = mlp_output + attn_output + hidden_states
+ else:
+ attn_output = attn_output + hidden_states
+ mlp_output = self.mlp(self.post_attention_layernorm(attn_output))
+ mlp_output = self.post_mlp_dropout(mlp_output)
+ hidden_states = mlp_output + attn_output
+
+ return hidden_states
+
+
+class GPTNeoXPreTrainedModel(LlamaPreTrainedModel):
+ base_model_prefix = "gpt_neox"
+ _no_split_modules = ["GPTNeoXLayer"]
+ _keys_to_ignore_on_load_unexpected = [r"attention.bias", r"attention.masked_bias"]
+ _can_record_outputs = {
+ "hidden_states": GPTNeoXLayer,
+ "attentions": GPTNeoXAttention,
+ }
+
+
+GPT_NEOX_START_DOCSTRING = None # Will be picked up by modular
+GPT_NEOX_INPUTS_DOCSTRING = None # Will be picked up by modular
+
+
+class GPTNeoXModel(LlamaModel):
+ def __init__(self, config):
+ PreTrainedModel.__init__(self, config)
+ self.config = config
+
+ self.embed_in = nn.Embedding(config.vocab_size, config.hidden_size)
+ self.emb_dropout = nn.Dropout(config.hidden_dropout)
+ self.layers = nn.ModuleList([GPTNeoXLayer(config, i) for i in range(config.num_hidden_layers)])
+ self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.rotary_emb = GPTNeoXRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embed_in
+
+ def set_input_embeddings(self, value):
+ self.embed_in = value
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_in(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ hidden_states = self.emb_dropout(inputs_embeds)
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+ for layer in self.layers:
+ hidden_states = layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_ids=position_ids,
+ layer_past=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = self.final_layer_norm(hidden_states)
+
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ GPTNeoX Model with a `language modeling` head on top for CLM fine-tuning.
+ """
+)
+class GPTNeoXForCausalLM(GPTNeoXPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"embed_out.weight": "gpt_neox.embed_in.weight"}
+ _tp_plan = {"embed_out": "colwise_gather_output"}
+ _pp_plan = {"embed_out": (["hidden_states"], ["logits"])}
+
+ def __init__(self, config):
+ super().__init__(config)
+
+ self.gpt_neox = GPTNeoXModel(config)
+ self.embed_out = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_output_embeddings(self):
+ return self.embed_out
+
+ def set_output_embeddings(self, new_embeddings):
+ self.embed_out = new_embeddings
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | CausalLMOutputWithPast:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the left-to-right language modeling loss (next word prediction). Indices should be in
+ `[-100, 0, ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are
+ ignored (masked), the loss is only computed for the tokens with labels n `[0, ..., config.vocab_size]`.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, GPTNeoXForCausalLM, GPTNeoXConfig
+ >>> import torch
+
+ >>> tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
+ >>> config = GPTNeoXConfig.from_pretrained("EleutherAI/gpt-neox-20b")
+ >>> config.is_decoder = True
+ >>> model = GPTNeoXForCausalLM.from_pretrained("EleutherAI/gpt-neox-20b", config=config)
+
+ >>> inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
+ >>> outputs = model(**inputs)
+
+ >>> prediction_logits = outputs.logits
+ ```"""
+
+ outputs: BaseModelOutputWithPast = self.gpt_neox(
+ input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.embed_out(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ The GPTNeoX Model transformer with a sequence classification head on top (linear layer).
+
+ [`GPTNeoXForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+ (e.g. GPT-1) do.
+
+ Since it does classification on the last token, it requires to know the position of the last token. If a
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+ each row of the batch).
+ """
+)
+class GPTNeoXForSequenceClassification(GPTNeoXPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.gpt_neox = GPTNeoXModel(config)
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ past_key_values: Cache | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs,
+ ) -> SequenceClassifierOutputWithPast:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+
+ outputs: BaseModelOutputWithPast = self.gpt_neox(
+ input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+ hidden_states = outputs.last_hidden_state
+ logits = self.score(hidden_states)
+
+ batch_size = logits.shape[0]
+ if self.config.pad_token_id is None and batch_size != 1:
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+ if self.config.pad_token_id is None:
+ last_non_pad_token = -1
+ elif input_ids is not None:
+ # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id
+ non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)
+ token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)
+ last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)
+ else:
+ last_non_pad_token = -1
+ logger.warning_once(
+ f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+ "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+ )
+
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)
+
+ return SequenceClassifierOutputWithPast(
+ loss=loss,
+ logits=pooled_logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+class GPTNeoXForTokenClassification(GPTNeoXPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+
+ self.gpt_neox = GPTNeoXModel(config)
+ self.dropout = nn.Dropout(config.classifier_dropout)
+ self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs,
+ ) -> TokenClassifierOutput:
+ r"""
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+ """
+
+ outputs: BaseModelOutputWithPast = self.gpt_neox(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ hidden_states = self.dropout(hidden_states)
+ logits = self.classifier(hidden_states)
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits, labels, self.config)
+
+ return TokenClassifierOutput(
+ loss=loss,
+ logits=logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+@auto_docstring
+class GPTNeoXForQuestionAnswering(GPTNeoXPreTrainedModel):
+ def __init__(self, config):
+ super().__init__(config)
+ self.num_labels = config.num_labels
+ self.gpt_neox = GPTNeoXModel(config)
+ self.qa_outputs = nn.Linear(config.hidden_size, 2)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.FloatTensor | None = None,
+ token_type_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ start_positions: torch.LongTensor | None = None,
+ end_positions: torch.LongTensor | None = None,
+ **kwargs,
+ ) -> QuestionAnsweringModelOutput:
+ outputs: BaseModelOutputWithPast = self.gpt_neox(
+ input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ inputs_embeds=inputs_embeds,
+ **kwargs,
+ )
+
+ sequence_output = outputs.last_hidden_state
+
+ logits = self.qa_outputs(sequence_output)
+ start_logits, end_logits = logits.split(1, dim=-1)
+ start_logits = start_logits.squeeze(-1).contiguous()
+ end_logits = end_logits.squeeze(-1).contiguous()
+
+ loss = None
+ if start_positions is not None and end_positions is not None:
+ loss = self.loss_function(start_logits, end_logits, start_positions, end_positions)
+
+ return QuestionAnsweringModelOutput(
+ loss=loss,
+ start_logits=start_logits,
+ end_logits=end_logits,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+__all__ = [
+ "GPTNeoXForCausalLM",
+ "GPTNeoXForQuestionAnswering",
+ "GPTNeoXForSequenceClassification",
+ "GPTNeoXForTokenClassification",
+ "GPTNeoXLayer",
+ "GPTNeoXModel",
+ "GPTNeoXPreTrainedModel",
+]
diff --git a/third_party/transformers/src/transformers/models/gpt_neox/tokenization_gpt_neox.py b/third_party/transformers/src/transformers/models/gpt_neox/tokenization_gpt_neox.py
new file mode 100644
index 0000000000000000000000000000000000000000..325c6225f49f507787395b18b36d2c2cb3f3bfe5
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/gpt_neox/tokenization_gpt_neox.py
@@ -0,0 +1,144 @@
+# Copyright 2022 EleutherAI and The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Tokenization classes for GPTNeoX."""
+
+from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers
+from tokenizers.models import BPE
+
+from ...tokenization_utils_tokenizers import TokenizersBackend
+from ...utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.json", "merges_file": "merges.txt", "tokenizer_file": "tokenizer.json"}
+
+
+class GPTNeoXTokenizer(TokenizersBackend):
+ """
+ Construct a GPT-NeoX-20B tokenizer (backed by HuggingFace's tokenizers library). Based on byte-level
+ Byte-Pair-Encoding.
+
+ This tokenizer has been trained to treat spaces like parts of the tokens (a bit like sentencepiece) so a word will
+ be encoded differently whether it is at the beginning of the sentence (without space) or not:
+
+ ```python
+ >>> from transformers import GPTNeoXTokenizer
+
+ >>> tokenizer = GPTNeoXTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
+ >>> tokenizer("Hello world")["input_ids"]
+ [15496, 995]
+
+ >>> tokenizer(" Hello world")["input_ids"]
+ [18435, 995]
+ ```
+
+ You can get around that behavior by passing `add_prefix_space=True` when instantiating this tokenizer, but since
+ the model was not pretrained this way, it might yield a decrease in performance.
+
+
+
+ When used with `is_split_into_words=True`, this tokenizer needs to be instantiated with `add_prefix_space=True`.
+
+
+
+ This tokenizer inherits from [`TokenizersBackend`] which contains most of the main methods. Users should
+ refer to this superclass for more information regarding those methods.
+
+ Args:
+ vocab_file (`str`, *optional*):
+ Path to the vocabulary file.
+ merges_file (`str`, *optional*):
+ Path to the merges file.
+ tokenizer_file (`str`, *optional*):
+ Path to a tokenizers JSON file containing the serialization of a tokenizer.
+ errors (`str`, *optional*, defaults to `"replace"`):
+ Paradigm to follow when decoding bytes to UTF-8. See
+ [bytes.decode](https://docs.python.org/3/library/stdtypes.html#bytes.decode) for more information.
+ unk_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+ The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
+ token instead.
+ bos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+ The beginning of sequence token.
+ eos_token (`str`, *optional*, defaults to `"<|endoftext|>"`):
+ The end of sequence token.
+ pad_token (`str`, *optional*, defaults to `"<|padding|>"`):
+ Token for padding a sequence.
+ add_prefix_space (`bool`, *optional*, defaults to `False`):
+ Whether or not to add an initial space to the input. This allows to treat the leading word just as any
+ other word. (GPTNeoX tokenizer detect beginning of words by the preceding space).
+ add_bos_token (`bool`, *optional*, defaults to `False`):
+ Whether or not to add a `bos_token` at the start of sequences.
+ add_eos_token (`bool`, *optional*, defaults to `False`):
+ Whether or not to add an `eos_token` at the end of sequences.
+ trim_offsets (`bool`, *optional*, defaults to `True`):
+ Whether or not the post-processing step should trim offsets to avoid including whitespaces.
+ vocab (`str` or `dict[str, int]`, *optional*):
+ Custom vocabulary dictionary. If not provided, vocabulary is loaded from `vocab_file`.
+ merges (`str` or `list[str]`, *optional*):
+ Custom merges list. If not provided, merges are loaded from `merges_file`.
+ """
+
+ vocab_files_names = VOCAB_FILES_NAMES
+ model_input_names = ["input_ids", "attention_mask"]
+ model = BPE
+
+ def __init__(
+ self,
+ vocab: str | dict[str, int] | None = None,
+ merges: str | list[str] | None = None,
+ errors: str = "replace",
+ unk_token: str = "<|endoftext|>",
+ bos_token: str = "<|endoftext|>",
+ eos_token: str = "<|endoftext|>",
+ pad_token: str = "<|padding|>",
+ add_prefix_space: bool = False,
+ trim_offsets: bool = True,
+ **kwargs,
+ ):
+ self.add_prefix_space = add_prefix_space
+ self.trim_offsets = trim_offsets
+
+ self._vocab = vocab if vocab is not None else {str(unk_token): 0, str(pad_token): 1}
+ self._merges = merges or []
+ self._tokenizer = Tokenizer(
+ BPE(
+ vocab=self._vocab,
+ merges=self._merges,
+ dropout=None,
+ continuing_subword_prefix="",
+ end_of_word_suffix="",
+ fuse_unk=False,
+ )
+ )
+
+ self._tokenizer.normalizer = normalizers.NFC()
+ self._tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(
+ add_prefix_space=add_prefix_space, trim_offsets=trim_offsets
+ )
+ self._tokenizer.decoder = decoders.ByteLevel()
+
+ super().__init__(
+ errors=errors,
+ unk_token=unk_token,
+ bos_token=bos_token,
+ eos_token=eos_token,
+ pad_token=pad_token,
+ add_prefix_space=add_prefix_space,
+ trim_offsets=trim_offsets,
+ **kwargs,
+ )
+
+
+__all__ = ["GPTNeoXTokenizer"]
diff --git a/third_party/transformers/src/transformers/models/granite_speech/__init__.py b/third_party/transformers/src/transformers/models/granite_speech/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..d6122581855250911b0bf951a42e22d9c354240a
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/granite_speech/__init__.py
@@ -0,0 +1,29 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_granite_speech import *
+ from .feature_extraction_granite_speech import *
+ from .modeling_granite_speech import *
+ from .processing_granite_speech import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/granite_speech/configuration_granite_speech.py b/third_party/transformers/src/transformers/models/granite_speech/configuration_granite_speech.py
new file mode 100644
index 0000000000000000000000000000000000000000..d02ac9998696d3cd2fed36a4459d653379d58ab7
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/granite_speech/configuration_granite_speech.py
@@ -0,0 +1,139 @@
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# 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.
+"""Config class for Granite Speech."""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+@auto_docstring(checkpoint="ibm-granite/granite-speech-3.3-2b")
+@strict
+class GraniteSpeechEncoderConfig(PreTrainedConfig):
+ r"""
+ feedforward_mult (`int`, *optional*, defaults to 4):
+ Multiplier for the up/down projections in the encoder's feedforward layers;
+ The projections will have intermediate dim of size `hidden_dim * feedforward_mult`.
+ output_dim (`int`, *optional*, defaults to 42):
+ Intermediate dimension of the feedforward projections in the conformer
+ to be added to every other encoder block's output.
+ context_size (`int`, *optional*, defaults to 200):
+ Context size to be used in conformer attention.
+ max_pos_emb (`int`, *optional*, defaults to 512):
+ Max pos embeds to be used in attention (shaw's relative positional encoding).
+ conv_expansion_factor (`int`, *optional*, defaults to 2):
+ Intermediate dimension to be used in conformer convolutions.
+
+ Example:
+
+ ```python
+ >>> from transformers import GraniteSpeechEncoderConfig, GraniteSpeechCTCEncoder
+
+ >>> # Initializing a GraniteSpeechEncoderConfig
+ >>> configuration = GraniteSpeechEncoderConfig()
+
+ >>> # Initializing a GraniteSpeechCTCEncoder (with random weights)
+ >>> model = GraniteSpeechCTCEncoder(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "granite_speech_encoder"
+
+ input_dim: int = 160
+ num_layers: int = 10
+ hidden_dim: int = 1024
+ feedforward_mult: int = 4
+ num_heads: int = 8
+ dim_head: int = 128
+ output_dim: int = 42
+ context_size: int = 200
+ max_pos_emb: int = 512
+ dropout: float | int = 0.1
+ conv_kernel_size: int = 15
+ conv_expansion_factor: int = 2
+
+
+@auto_docstring(checkpoint="ibm-granite/granite-speech-3.3-2b")
+@strict
+class GraniteSpeechConfig(PreTrainedConfig):
+ r"""
+ projector_config (`Union[AutoConfig, dict]`, *optional*, defaults to `Blip2QFormerConfig`):
+ The config object or dictionary of the audio projector.
+ has_lora_adapter (`bool`, *optional*, defaults to `True`):
+ Indicates whether or not the model has a lora adapter that should only
+ be activate when processing audio inputs.
+ downsample_rate (`int`, *optional*, defaults to 5):
+ Downsample rate for the audio feature extractor.
+ window_size (`int`, *optional*, defaults to 15):
+ Window size for the audio feature projector.
+
+ Example:
+
+ ```python
+ >>> from transformers import GraniteSpeechConfig, GraniteSpeechForConditionalGeneration
+
+ >>> # Initializing a GraniteSpeechConfig
+ >>> configuration = GraniteSpeechConfig()
+
+ >>> # Initializing a GraniteSpeechForConditionalGeneration (with random weights)
+ >>> model = GraniteSpeechForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "granite_speech"
+ attribute_map = {
+ "audio_token_id": "audio_token_index",
+ }
+ sub_configs = {
+ "text_config": AutoConfig,
+ "encoder_config": GraniteSpeechEncoderConfig,
+ "projector_config": AutoConfig,
+ }
+
+ text_config: dict | PreTrainedConfig | None = None
+ encoder_config: dict | PreTrainedConfig | None = None
+ projector_config: dict | PreTrainedConfig | None = None
+ audio_token_index: int = 49155
+ initializer_range: float = 0.02
+ has_lora_adapter: bool = True
+ downsample_rate: int = 5
+ window_size: int = 15
+
+ def __post_init__(self, **kwargs):
+ if isinstance(self.text_config, dict):
+ self.text_config["model_type"] = self.text_config.get("model_type", "granite")
+ self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config)
+ elif self.text_config is None:
+ self.text_config = CONFIG_MAPPING["granite"]()
+
+ if isinstance(self.projector_config, dict):
+ self.projector_config["model_type"] = self.projector_config.get("model_type", "blip_2_qformer")
+ self.projector_config = CONFIG_MAPPING[self.projector_config["model_type"]](**self.projector_config)
+ elif self.projector_config is None:
+ self.projector_config = CONFIG_MAPPING["blip_2_qformer"]()
+
+ if not isinstance(self.encoder_config, GraniteSpeechEncoderConfig):
+ self.encoder_config = {} if self.encoder_config is None else self.encoder_config
+ self.encoder_config = GraniteSpeechEncoderConfig(**self.encoder_config)
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["GraniteSpeechEncoderConfig", "GraniteSpeechConfig"]
diff --git a/third_party/transformers/src/transformers/models/granite_speech/feature_extraction_granite_speech.py b/third_party/transformers/src/transformers/models/granite_speech/feature_extraction_granite_speech.py
new file mode 100644
index 0000000000000000000000000000000000000000..cd32d0433bae4a8b34be30477876bbe648dba857
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/granite_speech/feature_extraction_granite_speech.py
@@ -0,0 +1,184 @@
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# 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.
+"""Feature extractor class for Granite Speech."""
+
+import math
+from collections.abc import Sequence
+
+import numpy as np
+
+from ...feature_extraction_utils import BatchFeature, FeatureExtractionMixin
+from ...tokenization_utils_base import AudioInput
+from ...utils import is_torch_available, is_torchaudio_available, logging
+from ...utils.import_utils import requires_backends
+
+
+logger = logging.get_logger(__name__)
+
+if is_torch_available():
+ import torch
+
+if is_torchaudio_available():
+ import torchaudio
+
+
+class GraniteSpeechFeatureExtractor(FeatureExtractionMixin):
+ model_input_names = ["input_features"]
+
+ def __init__(
+ self,
+ sampling_rate: int = 16000,
+ n_fft: int = 512,
+ win_length: int = 400,
+ hop_length: int = 160,
+ n_mels: int = 80,
+ projector_window_size: int = 15,
+ projector_downsample_rate: int = 5,
+ **kwargs,
+ ):
+ super().__init__(**kwargs)
+ self.sampling_rate = sampling_rate
+ self.melspec_kwargs = {
+ "sample_rate": sampling_rate,
+ "n_fft": n_fft,
+ "win_length": win_length,
+ "hop_length": hop_length,
+ "n_mels": n_mels,
+ }
+ requires_backends(self, ["torchaudio"])
+ self.mel_filters = torchaudio.transforms.MelSpectrogram(**self.melspec_kwargs)
+ self.projector_window_size = projector_window_size
+ self.projector_downsample_rate = projector_downsample_rate
+
+ def __call__(
+ self,
+ audios: AudioInput,
+ device: str | None = "cpu",
+ ) -> BatchFeature:
+ requires_backends(self, ["torchaudio"])
+
+ speech_inputs = {}
+ batched_audio, audio_lengths = self._get_audios_and_audio_lengths(audios)
+ speech_inputs["input_features"] = self._extract_mel_spectrograms(
+ batched_audio,
+ device=device,
+ )
+ audio_embed_sizes = self._get_num_audio_features(audio_lengths)
+ speech_inputs["audio_embed_sizes"] = audio_embed_sizes
+ # TODO (@alex-jw-brooks): Currently input_features_mask is not
+ # a great name, because input_features and input_features_mask
+ # have different shapes (before/after the projector).
+ #
+ # We should align this with other multimodal models, e.g,. llava
+ # and qwen2audio and refactor this to ensure input_feature_mask
+ # has the same dimensionality as input_features, or compute it in
+ # the model based on the audio embedding sizes (since we do not
+ # have an attention mask for the audio features to infer padding from).
+ speech_inputs["input_features_mask"] = torch.arange(max(audio_embed_sizes)).view(1, -1) < torch.tensor(
+ audio_embed_sizes
+ ).view(-1, 1)
+ return BatchFeature(data=speech_inputs)
+
+ def _extract_mel_spectrograms(self, audio: "torch.Tensor", device="cpu"):
+ """
+ Compute the Mel features to be passed to the conformer encoder.
+ """
+ requires_backends(self, ["torchaudio"])
+ if device is not None:
+ melspec = self.mel_filters.to(device)
+ audio = audio.to(device)
+ else:
+ melspec = self.mel_filters
+
+ bsz = audio.shape[0]
+ with torch.no_grad():
+ # Compute mel features
+ mel = melspec(audio.float())
+ logmel = mel.transpose(-1, -2).clip_(min=1e-10).log10_()
+ mx = logmel.amax(dim=(-2, -1), keepdim=True)
+ logmel = torch.maximum(logmel, mx - 8.0).div_(4).add_(1)
+ # remove last frame if odd
+ if logmel.shape[1] % 2 == 1:
+ logmel = logmel[:, :-1]
+
+ # stacking and skipping by 2
+ audio = logmel.reshape(bsz, -1, 2 * logmel.shape[-1])
+
+ return audio
+
+ def _get_num_audio_features(self, audio_lengths: Sequence[int]) -> Sequence[int]:
+ """
+ Gets the (variable length) number of features (i.e., projector output) for the sequences
+ being considered.
+
+ Args:
+ audio_lengths (`Sequence[int]`):
+ Sequence of one or more raw audio lengths.
+ """
+ hop_length = self.melspec_kwargs["hop_length"]
+ effective_window_size = self.projector_window_size // self.projector_downsample_rate
+
+ projector_lengths = []
+ for raw_length in audio_lengths:
+ # mel sequence length computation
+ mel_length = raw_length // hop_length + 1
+ # encoder frame takes two mel features
+ encoder_length = mel_length // 2
+ nblocks = math.ceil(encoder_length / self.projector_window_size)
+ # projector output length
+ projector_length = nblocks * effective_window_size
+ projector_lengths.append(projector_length)
+
+ return projector_lengths
+
+ def _get_audios_and_audio_lengths(self, audios: AudioInput) -> Sequence["torch.Tensor", Sequence[int]]:
+ """
+ Coerces audio inputs to torch tensors and extracts audio lengths prior to stacking.
+
+ Args:
+ audios (`AudioInput`):
+ Audio sequence, numpy array, or torch tensor.
+ """
+ requires_backends(self, ["torch"])
+
+ # Coerce to PyTorch tensors if we have numpy arrays, since
+ # currently we have a dependency on torch/torchaudio anyway
+ if isinstance(audios, np.ndarray):
+ audios = torch.from_numpy(audios)
+ elif isinstance(audios, Sequence) and isinstance(audios[0], np.ndarray):
+ audios = [torch.from_numpy(arr) for arr in audios]
+
+ if isinstance(audios, torch.Tensor):
+ if audios.ndim == 1:
+ audios = audios.unsqueeze(0)
+ if not torch.is_floating_point(audios):
+ raise ValueError("Invalid audio provided. Audio should be a floating point between 0 and 1")
+
+ if audios.shape[0] > 1:
+ logger.warning("Audio samples are already collated; assuming they all have the same length")
+ lengths = [audios.shape[-1]] * audios.shape[0]
+ return audios, lengths
+
+ elif isinstance(audios, Sequence) and isinstance(audios[0], torch.Tensor):
+ if not torch.is_floating_point(audios[0]):
+ raise ValueError("Invalid audio provided. Audio should be a floating point between 0 and 1")
+ lengths = [audio.shape[-1] for audio in audios]
+ audios = [audio.squeeze(0) for audio in audios]
+ audios = torch.nn.utils.rnn.pad_sequence(audios, batch_first=True, padding_value=0.0)
+ return audios, lengths
+
+ raise TypeError("Invalid audio provided. Audio should be a one or more torch tensors or numpy arrays")
+
+
+__all__ = ["GraniteSpeechFeatureExtractor"]
diff --git a/third_party/transformers/src/transformers/models/granite_speech/modeling_granite_speech.py b/third_party/transformers/src/transformers/models/granite_speech/modeling_granite_speech.py
new file mode 100644
index 0000000000000000000000000000000000000000..0fbc1d1035bf333e0ee1dca96086748c0788b89f
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/granite_speech/modeling_granite_speech.py
@@ -0,0 +1,587 @@
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# 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.
+
+import math
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ... import initialization as init
+from ...cache_utils import Cache
+from ...generation import GenerationMixin
+from ...modeling_outputs import BaseModelOutputWithPooling, ModelOutput
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import (
+ TransformersKwargs,
+ auto_docstring,
+ can_return_tuple,
+ is_peft_available,
+ logging,
+ torch_compilable_check,
+)
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from ..auto import AutoModel, AutoModelForCausalLM
+from .configuration_granite_speech import GraniteSpeechConfig, GraniteSpeechEncoderConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for LlavaNext causal language model (or autoregressive) outputs.
+ """
+)
+class GraniteSpeechCausalLMOutputWithPast(ModelOutput):
+ r"""
+ 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 (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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.
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+
+
+### Projector
+class GraniteSpeechEncoderProjector(nn.Module):
+ def __init__(self, config: GraniteSpeechConfig):
+ super().__init__()
+ self.hidden_size = config.projector_config.hidden_size
+ self.downsample_rate = config.downsample_rate
+ self.window_size = config.window_size
+ self.num_queries = config.window_size // config.downsample_rate
+
+ self.query = nn.Parameter(torch.zeros(1, self.num_queries, config.projector_config.hidden_size))
+ self.query.data.normal_(mean=0.0, std=1.0)
+
+ # By default, this will be a blip_2_qformer config
+ self.qformer = AutoModel.from_config(config.projector_config)
+ self.linear = nn.Linear(config.projector_config.hidden_size, config.text_config.hidden_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ batch_size, seq_len, dim = hidden_states.size()
+ nblocks = math.ceil(seq_len / self.window_size)
+ pad = nblocks * self.window_size - seq_len
+ hidden_states = nn.functional.pad(hidden_states, (0, 0, 0, pad), "constant", 0)
+ hidden_states = hidden_states.view(batch_size * nblocks, self.window_size, dim)
+
+ query_output = self.qformer(
+ query_embeds=self.query,
+ encoder_hidden_states=hidden_states,
+ encoder_attention_mask=None,
+ return_dict=True,
+ )
+ query_proj = self.linear(
+ query_output.last_hidden_state.view(batch_size, nblocks * self.window_size // self.downsample_rate, -1)
+ )
+ return query_proj
+
+
+### Encoder - conformer is adapted from: https://github.com/lucidrains/conformer.git
+class GraniteSpeechConformerFeedForward(nn.Module):
+ """Feedforward module for conformer encoder blocks."""
+
+ def __init__(self, config: GraniteSpeechEncoderConfig):
+ super().__init__()
+ self.pre_norm = nn.LayerNorm(config.hidden_dim)
+ self.up_proj = nn.Linear(config.hidden_dim, config.hidden_dim * config.feedforward_mult)
+ self.silu = nn.SiLU()
+ self.dropout = nn.Dropout(config.dropout)
+ self.down_proj = nn.Linear(config.hidden_dim * config.feedforward_mult, config.hidden_dim)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.pre_norm(hidden_states)
+ hidden_states = self.up_proj(hidden_states)
+ hidden_states = self.dropout(self.silu(hidden_states))
+ hidden_states = self.down_proj(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ return hidden_states
+
+
+class GraniteSpeechConformerAttention(nn.Module):
+ """Attention for conformer blocks using Shaw's relative positional embeddings.
+ See the following [paper](https://huggingface.co/papers/1803.02155) for more details.
+ """
+
+ def __init__(self, config: GraniteSpeechEncoderConfig):
+ super().__init__()
+
+ inner_dim = config.dim_head * config.num_heads
+ self.max_pos_emb = config.max_pos_emb
+ self.context_size = config.context_size
+ self.num_heads = config.num_heads
+ self.dim_head = config.dim_head
+ self.scale = self.dim_head**-0.5
+ self.pre_norm = nn.LayerNorm(config.hidden_dim)
+ self.to_q = nn.Linear(config.hidden_dim, inner_dim, bias=False)
+ self.to_kv = nn.Linear(config.hidden_dim, inner_dim * 2, bias=False)
+ self.to_out = nn.Linear(inner_dim, config.hidden_dim)
+ self.rel_pos_emb = nn.Embedding(2 * self.max_pos_emb + 1, self.dim_head)
+ self.dropout = nn.Dropout(config.dropout)
+
+ if self.context_size <= 0 or self.context_size > self.max_pos_emb:
+ raise ValueError("Context size is either less than 0 or exceeds the max_pos_emb")
+
+ def forward(self, hidden_states: torch.Tensor, attention_dists: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.pre_norm(hidden_states)
+ bsz, num_features, _ = hidden_states.shape
+
+ num_blocks = math.ceil(num_features / self.context_size)
+ remainder = num_features % self.context_size
+ if remainder > 0:
+ # right padding to reach block size
+ hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, self.context_size - remainder))
+
+ query_states = self.to_q(hidden_states)
+ key_states, value_states = self.to_kv(hidden_states).chunk(2, dim=-1)
+
+ query_states = query_states.reshape(bsz, num_blocks, self.context_size, self.num_heads, -1).transpose(2, 3)
+ key_states = key_states.reshape(bsz, num_blocks, self.context_size, self.num_heads, -1).transpose(2, 3)
+ value_states = value_states.reshape(bsz, num_blocks, self.context_size, self.num_heads, -1).transpose(2, 3)
+
+ # shaw's relative positional embedding
+ rel_pos_emb = self.rel_pos_emb(attention_dists)
+ # alternative computation of `pos_attn` - for readability
+ # rel_pos_emb_expanded = rel_pos_emb.view([1, 1, 1] + list(rel_pos_emb.shape))
+ # pos_attn = torch.sum(query_states.unsqueeze(-2) * rel_pos_emb_expanded, dim=-1) * self.scale
+ # einsum implementation of pos_attn - gives x30 speedup over the alternative
+ # TODO (@avihu111) find a fast alternative to einsum
+ pos_attn = torch.einsum("b m h c d, c r d -> b m h c r", query_states, rel_pos_emb) * self.scale
+
+ if remainder > 0:
+ # masked attention in the extended block
+ mask = torch.ones(self.context_size, self.context_size, dtype=bool, device=hidden_states.device)
+ mask[:remainder, :remainder] = 0
+ mask_value = -torch.finfo(pos_attn.dtype).max
+ pos_attn[:, -1, :].masked_fill_(mask, mask_value)
+
+ with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):
+ out = F.scaled_dot_product_attention(
+ query_states, key_states, value_states, attn_mask=pos_attn, scale=self.scale
+ )
+ out = out.transpose(2, 3).reshape(bsz, hidden_states.shape[1], -1)
+ out = self.to_out(out[:, :num_features, :])
+ return self.dropout(out)
+
+
+class GraniteSpeechConformerDepthWiseConv1d(nn.Module):
+ """Wrapper for padded 1D pointwise convolution."""
+
+ def __init__(self, chan_in: int, chan_out: int, kernel_size: int):
+ super().__init__()
+ # Padding for the 1D conv is symmetric or close (i.e., offset by one).
+ pad = kernel_size // 2
+ pad_offset = (kernel_size + 1) % 2
+ self.padding = (pad, pad - pad_offset)
+
+ self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, groups=chan_in, bias=False)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = F.pad(hidden_states, self.padding)
+ return self.conv(hidden_states)
+
+
+class GraniteSpeechConformerConvModule(nn.Module):
+ """Conformer conv module consisting of several 1D/depthwise 1D convolutional layers."""
+
+ def __init__(self, config: GraniteSpeechEncoderConfig):
+ super().__init__()
+ inner_dim = config.hidden_dim * config.conv_expansion_factor
+
+ self.norm = nn.LayerNorm(config.hidden_dim)
+ self.up_conv = nn.Conv1d(config.hidden_dim, inner_dim * 2, 1)
+ self.glu = nn.GLU(dim=1)
+ self.depth_conv = GraniteSpeechConformerDepthWiseConv1d(
+ inner_dim,
+ inner_dim,
+ kernel_size=config.conv_kernel_size,
+ )
+ self.silu = nn.SiLU()
+ self.batch_norm = nn.BatchNorm1d(inner_dim)
+ self.down_conv = nn.Conv1d(inner_dim, config.hidden_dim, 1)
+ self.dropout = nn.Dropout(config.dropout)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.norm(hidden_states)
+ hidden_states = self.up_conv(hidden_states.permute(0, 2, 1))
+ hidden_states = self.glu(hidden_states)
+ hidden_states = self.depth_conv(hidden_states)
+ hidden_states = self.silu(self.batch_norm(hidden_states))
+ hidden_states = self.down_conv(hidden_states).permute(0, 2, 1)
+ hidden_states = self.dropout(hidden_states)
+ return hidden_states
+
+
+class GraniteSpeechConformerBlock(nn.Module):
+ """Conformer block, consisting largely of linear layers, attention, and convolutional layers."""
+
+ def __init__(self, config: GraniteSpeechEncoderConfig):
+ super().__init__()
+ self.ff1 = GraniteSpeechConformerFeedForward(config)
+ self.attn = GraniteSpeechConformerAttention(config)
+ self.conv = GraniteSpeechConformerConvModule(config)
+ self.ff2 = GraniteSpeechConformerFeedForward(config)
+ self.post_norm = nn.LayerNorm(config.hidden_dim)
+
+ def forward(self, hidden_states: torch.Tensor, attention_dists: torch.Tensor) -> torch.Tensor:
+ hidden_states = 0.5 * self.ff1(hidden_states) + hidden_states
+ hidden_states = self.attn(hidden_states, attention_dists=attention_dists) + hidden_states
+ hidden_states = self.conv(hidden_states) + hidden_states
+ hidden_states = 0.5 * self.ff2(hidden_states) + hidden_states
+ hidden_states = self.post_norm(hidden_states)
+ return hidden_states
+
+
+@auto_docstring
+class GraniteSpeechPreTrainedModel(PreTrainedModel):
+ config: GraniteSpeechConfig
+ input_modalities = ("audio", "text")
+
+ _supports_flash_attn = False # `blip_2_qformer` dependency does not allow for this
+ _supports_sdpa = True
+
+ @torch.no_grad()
+ def _init_weights(self, module: nn.Module):
+ """Initialize the weights."""
+ super()._init_weights(module)
+ if isinstance(module, GraniteSpeechEncoderProjector):
+ init.normal_(module.query)
+ elif isinstance(module, GraniteSpeechCTCEncoder):
+ context_size = module.config.context_size
+ seq = torch.arange(context_size)
+ relpos_dist = seq.view(-1, 1) - seq.view(1, -1)
+ attention_dists = torch.clamp(relpos_dist, -context_size, context_size) + module.config.max_pos_emb
+ init.copy_(module.attention_dists, attention_dists)
+
+
+class GraniteSpeechCTCEncoder(GraniteSpeechPreTrainedModel):
+ config: GraniteSpeechEncoderConfig
+ input_modalities = "audio"
+ _can_record_outputs = {
+ "hidden_states": GraniteSpeechConformerBlock,
+ "attentions": GraniteSpeechConformerAttention,
+ }
+
+ def __init__(self, config: GraniteSpeechEncoderConfig):
+ super().__init__(config)
+
+ # Precompute clamped relative positional encoding distances
+ seq = torch.arange(config.context_size)
+ relpos_dist = seq.view(-1, 1) - seq.view(1, -1)
+ attention_dists = torch.clamp(relpos_dist, -config.context_size, config.context_size) + config.max_pos_emb
+ self.register_buffer("attention_dists", attention_dists, persistent=False)
+ self.input_linear = nn.Linear(config.input_dim, config.hidden_dim, bias=True)
+ self.layers = nn.ModuleList([GraniteSpeechConformerBlock(config) for _ in range(config.num_layers)])
+
+ self.out = nn.Linear(config.hidden_dim, config.output_dim, bias=True)
+ self.out_mid = nn.Linear(config.output_dim, config.hidden_dim, bias=True)
+ self.num_layers = config.num_layers
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ def forward(
+ self, hidden_states: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
+ ) -> tuple | BaseModelOutputWithPooling:
+ hidden_states = self.input_linear(hidden_states)
+ for idx, layer in enumerate(self.layers, start=1):
+ hidden_states = layer(hidden_states, attention_dists=self.attention_dists)
+
+ if idx == self.num_layers // 2:
+ hidden_states_mid = hidden_states.clone()
+ hidden_states_mid = self.out(hidden_states_mid)
+ hidden_states += self.out_mid(nn.Softmax(dim=-1)(hidden_states_mid))
+
+ return BaseModelOutputWithPooling(last_hidden_state=hidden_states)
+
+
+@auto_docstring(
+ custom_intro="""
+ The Granite Speech model, which consists of an audio encoder, projector, and language model.
+ """
+)
+class GraniteSpeechForConditionalGeneration(GraniteSpeechPreTrainedModel, GenerationMixin):
+ def __init__(self, config: GraniteSpeechConfig):
+ super().__init__(config)
+ # NOTE: It doesn't matter when we initialize from config, but we should be careful
+ # to make sure this does not pick up the adapter_config if in the future we use
+ # from_pretrained or something similar, since that should be set by the composite
+ # model; don't need to consider it twice
+ self.language_model = AutoModelForCausalLM.from_config(config.text_config)
+
+ self.encoder = GraniteSpeechCTCEncoder(config.encoder_config)
+ self.projector = GraniteSpeechEncoderProjector(config)
+
+ if config.has_lora_adapter and not is_peft_available():
+ logger.warning(
+ "Config indicates that a lora adapter should be present, but "
+ "peft is not installed; this will cause the model to perform "
+ "incorrectly when audio inputs are provided. Please install "
+ "peft and reload the model!"
+ )
+
+ self.post_init()
+
+ def set_decoder(self, decoder):
+ self.language_model.set_decoder(decoder)
+
+ def get_decoder(self):
+ return self.language_model.get_decoder()
+
+ def set_input_embeddings(self, value):
+ self.language_model.set_input_embeddings(value)
+
+ def set_output_embeddings(self, new_embeddings):
+ self.language_model.set_output_embeddings(new_embeddings)
+
+ def get_input_embeddings(self):
+ return self.language_model.get_input_embeddings()
+
+ def get_output_embeddings(self):
+ return self.language_model.get_output_embeddings()
+
+ @can_return_tuple
+ @auto_docstring
+ def get_audio_features(
+ self, input_features: torch.Tensor, **kwargs: Unpack[TransformersKwargs]
+ ) -> tuple | BaseModelOutputWithPooling:
+ audio_outputs = self.encoder(input_features, return_dict=True, **kwargs)
+ projected_embeds = self.projector(audio_outputs.last_hidden_state)
+ audio_outputs.pooler_output = projected_embeds
+
+ return audio_outputs
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ input_features: torch.FloatTensor | None = None,
+ input_features_mask: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **lm_kwargs,
+ ) -> tuple[torch.Tensor] | GraniteSpeechCausalLMOutputWithPast:
+ r"""
+ input_features_mask (`torch.Tensor`, *optional*):
+ Mask to be applied to audio features prior to scattering into the language embeddings.
+ 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]`.
+ """
+ # TODO (@alex-jw-brooks) add an example to this docstring once models are released
+ 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.return_dict
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if input_features is not None and inputs_embeds is not None:
+ raise ValueError(
+ "You cannot specify both input_features and inputs_embeds at the same time, and must specify either one"
+ )
+
+ if inputs_embeds is None:
+ # Get the base embeddings; set all audio tokens to 0 index
+ # to avoid out of vocabulary issues with the LLM embedding.
+ # Audio features will be masked into is_audio_idx indices later.
+ is_audio_idx = input_ids == self.config.audio_token_id
+ llm_input_ids = input_ids.clone()
+ llm_input_ids[is_audio_idx] = 0
+ inputs_embeds = self.get_input_embeddings()(llm_input_ids)
+
+ if input_features is not None:
+ if input_features.dtype != self.dtype:
+ input_features = input_features.to(self.dtype)
+ # Get the audio features from the encoder / projector
+ audio_embeds = self.get_audio_features(input_features, return_dict=True).pooler_output
+
+ # Merge the audio features into the LLM embeddings
+ inputs_embeds = self.get_merged_audio_embeddings(
+ input_ids=input_ids,
+ audio_features=audio_embeds,
+ input_features_mask=input_features_mask,
+ )
+
+ outputs = self.language_model(
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ logits_to_keep=logits_to_keep,
+ **lm_kwargs,
+ )
+ logits = outputs[0]
+
+ loss = None
+ if labels is not None:
+ # Shift so that tokens < n predict n
+ if attention_mask is not None:
+ # we use the input attention mask to shift the logits and labels, because it is 2D.
+ # we also crop attn mask in case it is longer, which happens in PrefixTuning with peft
+ shift_attention_mask = attention_mask[:, -(logits.shape[1] - 1) :].to(logits.device)
+ shift_logits = logits[..., :-1, :][shift_attention_mask.to(logits.device) != 0].contiguous()
+ shift_labels = labels[..., 1:][shift_attention_mask.to(labels.device) != 0].contiguous()
+ else:
+ shift_logits = logits[..., :-1, :].contiguous()
+ shift_labels = labels[..., 1:].contiguous()
+ # Flatten the tokens
+ loss_fct = nn.CrossEntropyLoss()
+ loss = loss_fct(
+ shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1).to(shift_logits.device)
+ )
+
+ if not return_dict:
+ output = (logits,) + outputs[1:]
+ return (loss,) + output if loss is not None else output
+
+ return GraniteSpeechCausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+ def prepare_inputs_for_generation(
+ self,
+ input_ids,
+ past_key_values=None,
+ inputs_embeds=None,
+ input_features=None,
+ attention_mask=None,
+ logits_to_keep=None,
+ is_first_iteration=False,
+ **kwargs,
+ ):
+ # Overwritten -- in specific circumstances we don't want to forward audio inputs to the model
+
+ model_inputs = self.language_model.prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ logits_to_keep=logits_to_keep,
+ is_first_iteration=is_first_iteration,
+ **kwargs,
+ )
+
+ # If we're in cached decoding stage, input_features should be None because
+ # input ids do not contain special audio token anymore Otherwise we need
+ # input feature values to be passed to the model
+ if is_first_iteration or not kwargs.get("use_cache", True):
+ model_inputs["input_features"] = input_features
+ return model_inputs
+
+ def get_merged_audio_embeddings(
+ self, input_ids: torch.Tensor, audio_features: torch.Tensor, input_features_mask: torch.Tensor | None = None
+ ) -> torch.Tensor:
+ """
+ Adds the audio token to the model's LLM vocabulary so that we can pass it
+ through the tokenizer; it's assumed that the embeddings corresponding to the
+ <|audio|> token will be clobbered with speech features.
+
+ Args:
+ input_ids (`torch.Tensor`):
+ Input IDs containing one or more audio tokens.
+ audio_features (`torch.Tensor`):
+ Audio features to be masked into the language embeddings to form multimodal embeddings.
+ input_features_mask (`torch.Tensor`, *optional*, defaults to `None`)
+ Mask to be applied to audio features prior to scattering into the language embeddings.
+ """
+ is_audio_index = input_ids == self.config.audio_token_id
+ llm_input_ids = torch.where(is_audio_index, 0, input_ids)
+ inputs_embeds = self.language_model.get_input_embeddings()(llm_input_ids) # [bsz, # features, hidden size]
+
+ # Mask the audio features into the text embeddings
+ special_audio_mask = is_audio_index.unsqueeze(-1)
+ audio_features = audio_features.to(inputs_embeds.device, inputs_embeds.dtype)
+ if input_features_mask is not None:
+ torch_compilable_check(
+ not torch.all(is_audio_index.int().sum(dim=1) != input_features_mask.int().sum(dim=1)),
+ "Number of audio tokens does not match number of audio features",
+ )
+ audio_features = audio_features[input_features_mask]
+
+ inputs_embeds = inputs_embeds.masked_scatter(
+ special_audio_mask,
+ audio_features,
+ )
+ return inputs_embeds
+
+ def generate(self, *args, **kwargs) -> torch.LongTensor:
+ # This model is expected to have a lora adapter, which is only
+ # enabled when considering audio inputs. As such, we override generate
+ # to conditionally enable / disable the lora adapter based on whether
+ # or not any input features were provided.
+
+ input_features = kwargs.pop("input_features", None)
+ if is_peft_available and self._hf_peft_config_loaded:
+ if input_features is not None:
+ self.enable_adapters()
+ else:
+ self.disable_adapters()
+ return super().generate(*args, input_features=input_features, **kwargs)
+
+ def save_pretrained(self, save_directory, *args, **kwargs):
+ # overwrite save_pretrained to first save the adapter if we have one
+ if is_peft_available and self._hf_peft_config_loaded:
+ adapter_name = self._get_adapter_name()
+ self.peft_config[adapter_name].base_model_name_or_path = save_directory
+ super().save_pretrained(save_directory, *args, **kwargs)
+ # Then save the base model afterwards
+ prev_val = self._hf_peft_config_loaded
+ self._hf_peft_config_loaded = False
+ super().save_pretrained(save_directory, *args, **kwargs)
+ self._hf_peft_config_loaded = prev_val
+
+ def _get_adapter_name(self):
+ return list(self.peft_config.keys())[0]
+
+
+__all__ = [
+ "GraniteSpeechCTCEncoder",
+ "GraniteSpeechForConditionalGeneration",
+ "GraniteSpeechPreTrainedModel",
+]
diff --git a/third_party/transformers/src/transformers/models/granite_speech/processing_granite_speech.py b/third_party/transformers/src/transformers/models/granite_speech/processing_granite_speech.py
new file mode 100644
index 0000000000000000000000000000000000000000..969ebc1c35965dd1ec90ae31027fd7afcabf7ef8
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/granite_speech/processing_granite_speech.py
@@ -0,0 +1,105 @@
+# Copyright 2025 The HuggingFace Inc. team.
+#
+# 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.
+"""Processor class for Granite Speech."""
+
+from typing import Union
+
+from ...feature_extraction_utils import BatchFeature
+from ...processing_utils import ProcessorMixin
+from ...tokenization_python import PreTokenizedInput, TextInput
+from ...utils import auto_docstring, is_torch_available, logging
+from ...utils.import_utils import requires_backends
+
+
+if is_torch_available():
+ import torch
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring
+class GraniteSpeechProcessor(ProcessorMixin):
+ def __init__(
+ self,
+ audio_processor,
+ tokenizer,
+ audio_token="<|audio|>",
+ chat_template=None,
+ ):
+ r"""
+ audio_token (`str`, *optional*, defaults to `"<|audio|>"`):
+ The special token used to represent audio in the text sequence. This token serves as a placeholder
+ that will be replaced with multiple audio tokens based on the actual audio length. The number of
+ audio tokens inserted depends on the audio feature dimensions extracted by the audio processor.
+ """
+ self.audio_token = tokenizer.audio_token if hasattr(tokenizer, "audio_token") else audio_token
+ super().__init__(audio_processor, tokenizer, chat_template=chat_template)
+
+ @auto_docstring
+ def __call__(
+ self,
+ text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput],
+ audio: Union["torch.Tensor", list["torch.Tensor"]] = None,
+ device: str = "cpu",
+ **kwargs,
+ ) -> BatchFeature:
+ requires_backends(self, ["torch"])
+
+ text = self._get_validated_text(text)
+ prompt_strings = text
+
+ if audio is not None:
+ # NOTE - we intentionally avoid throwing for potentially misaligned
+ # text / audio inputs here because some inference engines will
+ # trigger the conditions due to the way they call multimodal
+ # processors, e.g., vLLM.
+ audio_inputs = self.audio_processor(audio, device=device)
+
+ # TODO (@alex-jw-brooks); we should add a util to get_num_audio_tokens
+ # from feature lengths and call it here, rather than returning it
+ # from the feature extractor.
+ audio_embed_sizes = audio_inputs.pop("audio_embed_sizes")
+
+ # Expand the audio placeholders to match the feature dims; this
+ # is similar to how many VLMs handle image tokens, e.g., llava next
+ prompt_strings = []
+ num_replaced = 0
+ for sample in text:
+ while self.audio_token in sample:
+ sample = sample.replace(
+ self.audio_token,
+ "" * audio_embed_sizes[num_replaced],
+ 1,
+ )
+ num_replaced += 1
+ prompt_strings.append(sample)
+
+ prompt_strings = [sample.replace("", self.audio_token) for sample in prompt_strings]
+ else:
+ audio_inputs = {}
+
+ if "padding" not in kwargs:
+ kwargs["padding"] = True
+ text_inputs = self.tokenizer(prompt_strings, **kwargs)
+ return BatchFeature(data={**text_inputs, **audio_inputs})
+
+ def _get_validated_text(self, text: str | list) -> list[str]:
+ if isinstance(text, str):
+ return [text]
+ elif isinstance(text, list) and isinstance(text[0], str):
+ return text
+ raise TypeError("Invalid text provided! Text should be a string or list of strings.")
+
+
+__all__ = ["GraniteSpeechProcessor"]
diff --git a/third_party/transformers/src/transformers/models/granitemoeshared/__init__.py b/third_party/transformers/src/transformers/models/granitemoeshared/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..33d80cdd3425f95de5d40c82a4f52132be971f1f
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/granitemoeshared/__init__.py
@@ -0,0 +1,27 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_granitemoeshared import *
+ from .modeling_granitemoeshared import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/granitemoeshared/configuration_granitemoeshared.py b/third_party/transformers/src/transformers/models/granitemoeshared/configuration_granitemoeshared.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d782c089e3699ec67372a622f2dafc371aee617
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/granitemoeshared/configuration_granitemoeshared.py
@@ -0,0 +1,95 @@
+# Copyright 2024 EleutherAI 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.
+"""GraniteMoeShared model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="ibm-granite/granite-speech-3.2-8b")
+@strict
+class GraniteMoeSharedConfig(PreTrainedConfig):
+ r"""
+ embedding_multiplier (`float`, *optional*, defaults to 1.0):
+ embedding multiplier
+ logits_scaling (`float`, *optional*, defaults to 1.0):
+ divisor for output logits
+ residual_multiplier (`float`, *optional*, defaults to 1.0):
+ residual multiplier
+ attention_multiplier (`float`, *optional*, defaults to 1.0):
+ attention multiplier
+ shared_intermediate_size (`int`, *optional*, defaults to 1024):
+ intermediate size for shared experts.
+
+ ```python
+ >>> from transformers import GraniteMoeSharedModel, GraniteMoeSharedConfig
+
+ >>> # Initializing a GraniteMoeShared granitemoe-3b style configuration
+ >>> configuration = GraniteMoeSharedConfig()
+
+ >>> # Initializing a model from the granitemoe-7b style configuration
+ >>> model = GraniteMoeSharedModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```
+ """
+
+ model_type = "granitemoeshared"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ vocab_size: int = 32000
+ hidden_size: int = 4096
+ intermediate_size: int = 11008
+ num_hidden_layers: int = 32
+ num_attention_heads: int = 32
+ num_key_value_heads: int | None = None
+ hidden_act: str = "silu"
+ max_position_embeddings: int = 2048
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-6
+ use_cache: bool = True
+ pad_token_id: int | None = None
+ bos_token_id: int | None = 1
+ eos_token_id: int | list[int] | None = 2
+ tie_word_embeddings: bool = False
+ rope_parameters: RopeParameters | dict | None = None
+ attention_bias: bool = False
+ attention_dropout: float | int | None = 0.0
+ embedding_multiplier: float | int | None = 1.0
+ logits_scaling: float | int | None = 1.0
+ residual_multiplier: float | int | None = 1.0
+ attention_multiplier: float | int | None = 1.0
+ num_local_experts: int | None = 8
+ num_experts_per_tok: int | None = 2
+ output_router_logits: bool | None = False
+ router_aux_loss_coef: float | None = 0.001
+ shared_intermediate_size: int = 0
+
+ def __post_init__(self, **kwargs):
+ if self.num_key_value_heads is None:
+ self.num_key_value_heads = self.num_attention_heads
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["GraniteMoeSharedConfig"]
diff --git a/third_party/transformers/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py b/third_party/transformers/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py
new file mode 100644
index 0000000000000000000000000000000000000000..71f8c6eaff7dc0836d997ff844f8623b1aaa20c6
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/granitemoeshared/modeling_granitemoeshared.py
@@ -0,0 +1,800 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/granitemoeshared/modular_granitemoeshared.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_granitemoeshared.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2024 IBM and the HuggingFace Inc. team. All rights reserved.
+#
+#
+# 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.
+from collections.abc import Callable
+from typing import Optional, TypedDict
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
+from ...masking_utils import create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring
+from ...utils.generic import can_return_tuple, maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_granitemoeshared import GraniteMoeSharedConfig
+
+
+class GraniteFlashAttentionKwargs(TypedDict, total=False):
+ """
+ Keyword arguments for advanced Flash Attention, causal-conv1d, and mamba_ssm kernel usage.
+ Use cases include padding-free training and fewer `torch.compile` graph breaks.
+
+ cu_seq_lens_q (`torch.LongTensor`):
+ Gets cumulative sequence length for query state.
+ cu_seq_lens_k (`torch.LongTensor`):
+ Gets cumulative sequence length for key state.
+ max_length_q (`int`):
+ Maximum sequence length for query state.
+ max_length_k (`int`):
+ Maximum sequence length for key state.
+ seq_idx (`torch.IntTensor):
+ Index of each packed sequence.
+ """
+
+ cu_seq_lens_q: torch.LongTensor
+ cu_seq_lens_k: torch.LongTensor
+ max_length_q: int
+ max_length_k: int
+ seq_idx: torch.IntTensor
+
+
+class GraniteMoeSharedMLP(nn.Module):
+ """
+ MLP layer for shared experts
+
+ Args:
+ config:
+ Configuration object with model hyperparameters.
+ """
+
+ def __init__(self, config: GraniteMoeSharedConfig):
+ super().__init__()
+
+ self.input_size = config.hidden_size
+ self.hidden_size = config.shared_intermediate_size
+ self.activation = ACT2FN[config.hidden_act]
+ self.input_linear = nn.Linear(self.input_size, self.hidden_size * 2, bias=False)
+ self.output_linear = nn.Linear(self.hidden_size, self.input_size, bias=False)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.input_linear(hidden_states)
+ chunked_hidden_states = hidden_states.chunk(2, dim=-1)
+ hidden_states = self.activation(chunked_hidden_states[0]) * chunked_hidden_states[1]
+ hidden_states = self.output_linear(hidden_states)
+ return hidden_states
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class GraniteMoeSharedRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ GraniteMoeSharedRMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class GraniteMoeSharedParallelExperts(nn.Module):
+ def __init__(self, num_experts: int, input_size: int, output_size: int) -> None:
+ """
+ Initialize the GraniteMoeSharedParallelExperts module.
+ The experts weights are stored in [num_experts, output_size, input_size] format. Such that it's compatible with
+ many MoE libraries, such as [Megablock](https://github.com/databricks/megablocks) and
+ [ScatterMoE](https://github.com/shawntan/scattermoe), as well as the
+ [MoE kernel](https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/fused_moe/fused_moe.py)
+ used in vllm.
+
+ Args:
+ num_experts (int):
+ Number of experts.
+ input_size (int):
+ Size of the input.
+ output_size (int):
+ Size of the output.
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.empty(num_experts, output_size, input_size))
+ self.num_experts = num_experts
+ self.input_size = input_size
+ self.output_size = output_size
+
+ def forward(self, inputs, expert_size):
+ """
+ Forward pass of the GraniteMoeSharedParallelExperts module.
+
+ Args:
+ inputs (Tensor):
+ Input tensor.
+ expert_size:
+ Expert size information.
+
+ Returns:
+ Tensor: Output tensor.
+ """
+ input_list = inputs.split(expert_size, dim=0)
+ output_list = []
+ for i in range(self.num_experts):
+ output_list.append(F.linear(input_list[i], self.weight[i]))
+ results = torch.cat(output_list, dim=0)
+ return results
+
+
+class GraniteMoeSharedTopKGating(nn.Module):
+ def __init__(self, input_size: int, num_experts: int, top_k: int):
+ """
+ Initialize the top-k gating mechanism.
+
+ Args:
+ input_size (`int`):
+ Size of the input.
+ num_experts (`int`):
+ Number of experts.
+ top_k (`int`):
+ Number of top experts to select.
+ """
+ super().__init__()
+
+ self.num_experts = num_experts
+ self.input_size = input_size
+ self.top_k = top_k
+
+ self.layer = nn.Linear(input_size, num_experts, bias=False)
+
+ def forward(self, hidden_states):
+ # compute the top_k routing decision
+ logits = self.layer(hidden_states).float() # [batch_size x seq_len, num_experts]
+ top_k_logits, top_k_indices = logits.topk(self.top_k, dim=1) # [num_tokens, top_k]
+ top_k_gates = torch.softmax(top_k_logits, dim=1).type_as(hidden_states) # [num_tokens, top_k]
+
+ # compute number of input given to each expert
+ zeros = torch.zeros(
+ [top_k_gates.size(0), self.num_experts], dtype=top_k_gates.dtype, device=top_k_gates.device
+ ) # [num_tokens, num_experts]
+ gates = zeros.scatter(1, top_k_indices, 1) # [num_tokens, num_experts]
+ expert_size = gates.long().sum(0) # [num_experts,]
+ # (This cause torch.compile to fail with `torch._dynamo.exc.Unsupported: Backend compiler failed with a fake tensor exception at`)
+ # (and `DataDependentOutputException`)
+ expert_size = expert_size.tolist()
+
+ # sort and group input tokens according to expert assignment
+ top_k_experts = top_k_indices.flatten() # [num_tokens * top_k]
+ _, index_sorted_experts = top_k_experts.sort(0) # [num_tokens * top_k]
+ batch_index = index_sorted_experts.div(self.top_k, rounding_mode="trunc") # [num_tokens * top_k]
+
+ # gather the gate values for grouped input tokens
+ top_k_gates = top_k_gates.flatten() # [num_tokens * top_k]
+ batch_gates = top_k_gates[index_sorted_experts] # [num_tokens * top_k]
+
+ return index_sorted_experts, batch_index, batch_gates, expert_size, logits
+
+
+class GraniteMoeSharedMoE(nn.Module):
+ """
+ A Sparsely gated mixture of experts layer with 1-layer Feed-Forward networks as experts.
+
+ Args:
+ config:
+ Configuration object with model hyperparameters.
+ """
+
+ def __init__(self, config: GraniteMoeSharedConfig):
+ super().__init__()
+
+ self.input_size = config.hidden_size
+ self.hidden_size = config.intermediate_size
+ self.activation = ACT2FN[config.hidden_act]
+ self.input_linear = GraniteMoeSharedParallelExperts(
+ config.num_local_experts, self.input_size, self.hidden_size * 2
+ )
+ self.output_linear = GraniteMoeSharedParallelExperts(
+ config.num_local_experts, self.hidden_size, self.input_size
+ )
+
+ self.router = GraniteMoeSharedTopKGating(
+ input_size=self.input_size,
+ num_experts=config.num_local_experts,
+ top_k=config.num_experts_per_tok,
+ )
+
+ def forward(self, layer_input):
+ bsz, length, emb_size = layer_input.size()
+ layer_input = layer_input.reshape(-1, emb_size)
+ _, batch_index, batch_gates, expert_size, _ = self.router(layer_input)
+
+ expert_inputs = layer_input[batch_index]
+ hidden_states = self.input_linear(expert_inputs, expert_size)
+ chunked_hidden_states = hidden_states.chunk(2, dim=-1)
+ hidden_states = self.activation(chunked_hidden_states[0]) * chunked_hidden_states[1]
+ expert_outputs = self.output_linear(hidden_states, expert_size)
+
+ expert_outputs = expert_outputs * batch_gates[:, None]
+
+ zeros = torch.zeros((bsz * length, self.input_size), dtype=expert_outputs.dtype, device=expert_outputs.device)
+ layer_output = zeros.index_add(0, batch_index, expert_outputs)
+ layer_output = layer_output.view(bsz, length, self.input_size)
+ return layer_output
+
+
+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)
+
+
+@use_kernel_func_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+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)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class GraniteMoeSharedAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: GraniteMoeSharedConfig, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = config.attention_multiplier # Only diff with llama
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = True
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ attention_mask: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class GraniteMoeSharedDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: GraniteMoeSharedConfig, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = GraniteMoeSharedAttention(config=config, layer_idx=layer_idx)
+ self.input_layernorm = GraniteMoeSharedRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = GraniteMoeSharedRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.block_sparse_moe = GraniteMoeSharedMoE(config)
+ self.residual_multiplier = config.residual_multiplier # Only diff with mixtral!
+ self.shared_mlp = None if config.shared_intermediate_size == 0 else GraniteMoeSharedMLP(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ output_attentions: bool | None = False,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[GraniteFlashAttentionKwargs],
+ ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ output_attentions=output_attentions,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = residual + hidden_states * self.residual_multiplier
+
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ moe_hidden_states = self.block_sparse_moe(hidden_states)
+
+ if self.shared_mlp is None:
+ hidden_states = moe_hidden_states
+ else:
+ hidden_states = moe_hidden_states + self.shared_mlp(hidden_states)
+ hidden_states = residual + hidden_states * self.residual_multiplier
+ return hidden_states
+
+
+@auto_docstring
+class GraniteMoeSharedPreTrainedModel(PreTrainedModel):
+ config: GraniteMoeSharedConfig
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["GraniteMoeSharedDecoderLayer"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+ _can_compile_fullgraph = False # TopK gating fails fullgraph compilation at "expert_size = expert_size.tolist()"
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": GraniteMoeSharedDecoderLayer,
+ "attentions": GraniteMoeSharedAttention,
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ super()._init_weights(module)
+ if isinstance(module, GraniteMoeSharedParallelExperts):
+ init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+
+
+class GraniteMoeSharedRotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: GraniteMoeSharedConfig, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: GraniteMoeSharedConfig | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@auto_docstring
+class GraniteMoeSharedModel(GraniteMoeSharedPreTrainedModel):
+ def __init__(self, config: GraniteMoeSharedConfig):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList(
+ [GraniteMoeSharedDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = GraniteMoeSharedRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = GraniteMoeSharedRotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+ self.embedding_multiplier = config.embedding_multiplier
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> MoeModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask( # ONLY DIFF WITH MIXTRAL: NO SLIDING
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+ inputs_embeds = inputs_embeds * self.embedding_multiplier
+ hidden_states = inputs_embeds
+
+ # create position embeddings to be shared across the decoder layers
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
+
+ for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+ hidden_states = decoder_layer(
+ hidden_states,
+ position_embeddings=position_embeddings,
+ attention_mask=causal_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+
+ return MoeModelOutputWithPast( # only diff with Mistral is the output type, we need MoE
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+def load_balancing_loss_func(
+ gate_logits: torch.Tensor | tuple[torch.Tensor] | None,
+ num_experts: int | None = None,
+ top_k=2,
+ attention_mask: torch.Tensor | None = None,
+) -> torch.Tensor | int:
+ r"""
+ Computes auxiliary load balancing loss as in Switch Transformer - implemented in Pytorch.
+
+ See Switch Transformer (https://huggingface.co/papers/2101.03961) for more details. This function implements the loss
+ function presented in equations (4) - (6) of the paper. It aims at penalizing cases where the routing between
+ experts is too unbalanced.
+
+ Args:
+ gate_logits:
+ Logits from the `gate`, should be a tuple of model.config.num_hidden_layers tensors of
+ shape [batch_size X sequence_length, num_experts].
+ num_experts:
+ Number of experts
+ top_k:
+ The number of experts to route per-token, can be also interpreted as the `top-k` routing
+ parameter.
+ attention_mask (`torch.Tensor`, *optional*):
+ The attention_mask used in forward function
+ shape [batch_size X sequence_length] if not None.
+
+ Returns:
+ The auxiliary loss.
+ """
+ if gate_logits is None or not isinstance(gate_logits, tuple):
+ return 0
+
+ if isinstance(gate_logits, tuple):
+ compute_device = gate_logits[0].device
+ concatenated_gate_logits = torch.cat([layer_gate.to(compute_device) for layer_gate in gate_logits], dim=0)
+
+ routing_weights = torch.nn.functional.softmax(concatenated_gate_logits, dim=-1)
+
+ _, selected_experts = torch.topk(routing_weights, top_k, dim=-1)
+
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_experts)
+
+ if attention_mask is None:
+ # Compute the percentage of tokens routed to each experts
+ tokens_per_expert = torch.mean(expert_mask.float(), dim=0)
+
+ # Compute the average probability of routing to these experts
+ router_prob_per_expert = torch.mean(routing_weights, dim=0)
+ else:
+ batch_size, sequence_length = attention_mask.shape
+ num_hidden_layers = concatenated_gate_logits.shape[0] // (batch_size * sequence_length)
+
+ # Compute the mask that masks all padding tokens as 0 with the same shape of expert_mask
+ expert_attention_mask = (
+ attention_mask[None, :, :, None, None]
+ .expand((num_hidden_layers, batch_size, sequence_length, top_k, num_experts))
+ .reshape(-1, top_k, num_experts)
+ .to(compute_device)
+ )
+
+ # Compute the percentage of tokens routed to each experts
+ tokens_per_expert = torch.sum(expert_mask.float() * expert_attention_mask, dim=0) / torch.sum(
+ expert_attention_mask, dim=0
+ )
+
+ # Compute the mask that masks all padding tokens as 0 with the same shape of tokens_per_expert
+ router_per_expert_attention_mask = (
+ attention_mask[None, :, :, None]
+ .expand((num_hidden_layers, batch_size, sequence_length, num_experts))
+ .reshape(-1, num_experts)
+ .to(compute_device)
+ )
+
+ # Compute the average probability of routing to these experts
+ router_prob_per_expert = torch.sum(routing_weights * router_per_expert_attention_mask, dim=0) / torch.sum(
+ router_per_expert_attention_mask, dim=0
+ )
+
+ overall_loss = torch.sum(tokens_per_expert * router_prob_per_expert.unsqueeze(0))
+ return overall_loss * num_experts
+
+
+@auto_docstring
+class GraniteMoeSharedForCausalLM(GraniteMoeSharedPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+ _tp_plan = {"lm_head": "colwise_gather_output"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+ def __init__(self, config: GraniteMoeSharedConfig):
+ super().__init__(config)
+ self.model = GraniteMoeSharedModel(config)
+ self.vocab_size = config.vocab_size
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+ self.router_aux_loss_coef = config.router_aux_loss_coef
+ self.num_experts = config.num_local_experts
+ self.num_experts_per_tok = config.num_experts_per_tok
+ self.logits_scaling = config.logits_scaling
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @auto_docstring
+ @can_return_tuple
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ output_router_logits: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs,
+ ) -> tuple | MoeCausalLMOutputWithPast:
+ r"""
+ 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]`.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, GraniteMoeSharedForCausalLM
+
+ >>> model = GraniteMoeSharedForCausalLM.from_pretrained("ibm/PowerMoE-3b")
+ >>> tokenizer = AutoTokenizer.from_pretrained("ibm/PowerMoE-3b")
+
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+ ```"""
+ output_router_logits = (
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
+ )
+ # 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,
+ **kwargs,
+ )
+
+ # Only compute necessary logits
+ hidden_states = outputs.last_hidden_state
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+ logits = logits / self.config.logits_scaling
+
+ loss = None
+ if labels is not None:
+ # Flatten the tokens
+ loss = self.loss_function(
+ logits,
+ labels,
+ vocab_size=self.config.vocab_size,
+ **kwargs,
+ )
+
+ aux_loss = None
+ if output_router_logits:
+ aux_loss = load_balancing_loss_func(
+ outputs.router_logits,
+ self.num_experts,
+ self.num_experts_per_tok,
+ attention_mask,
+ )
+ if labels is not None:
+ loss += self.router_aux_loss_coef * aux_loss.to(loss.device) # make sure to reside in the same device
+ return MoeCausalLMOutputWithPast(
+ loss=loss,
+ aux_loss=aux_loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ router_logits=outputs.router_logits,
+ )
+
+
+__all__ = ["GraniteMoeSharedForCausalLM", "GraniteMoeSharedModel", "GraniteMoeSharedPreTrainedModel"]
diff --git a/third_party/transformers/src/transformers/models/granitemoeshared/modular_granitemoeshared.py b/third_party/transformers/src/transformers/models/granitemoeshared/modular_granitemoeshared.py
new file mode 100644
index 0000000000000000000000000000000000000000..e51cd7712b9b2ffacf2b3493056636e7aeb23a83
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/granitemoeshared/modular_granitemoeshared.py
@@ -0,0 +1,154 @@
+# Copyright 2024 IBM and the HuggingFace Inc. team. All rights reserved.
+#
+#
+# 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.
+from typing import TypedDict
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...cache_utils import Cache
+from ...processing_utils import Unpack
+from ...utils import logging
+from ..granitemoe.modeling_granitemoe import (
+ GraniteMoeDecoderLayer,
+ GraniteMoeForCausalLM,
+ GraniteMoeModel,
+ GraniteMoePreTrainedModel,
+)
+from .configuration_granitemoeshared import GraniteMoeSharedConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+class GraniteFlashAttentionKwargs(TypedDict, total=False):
+ """
+ Keyword arguments for advanced Flash Attention, causal-conv1d, and mamba_ssm kernel usage.
+ Use cases include padding-free training and fewer `torch.compile` graph breaks.
+
+ cu_seq_lens_q (`torch.LongTensor`):
+ Gets cumulative sequence length for query state.
+ cu_seq_lens_k (`torch.LongTensor`):
+ Gets cumulative sequence length for key state.
+ max_length_q (`int`):
+ Maximum sequence length for query state.
+ max_length_k (`int`):
+ Maximum sequence length for key state.
+ seq_idx (`torch.IntTensor):
+ Index of each packed sequence.
+ """
+
+ cu_seq_lens_q: torch.LongTensor
+ cu_seq_lens_k: torch.LongTensor
+ max_length_q: int
+ max_length_k: int
+ seq_idx: torch.IntTensor
+
+
+class GraniteMoeSharedMLP(nn.Module):
+ """
+ MLP layer for shared experts
+
+ Args:
+ config:
+ Configuration object with model hyperparameters.
+ """
+
+ def __init__(self, config: GraniteMoeSharedConfig):
+ super().__init__()
+
+ self.input_size = config.hidden_size
+ self.hidden_size = config.shared_intermediate_size
+ self.activation = ACT2FN[config.hidden_act]
+ self.input_linear = nn.Linear(self.input_size, self.hidden_size * 2, bias=False)
+ self.output_linear = nn.Linear(self.hidden_size, self.input_size, bias=False)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.input_linear(hidden_states)
+ chunked_hidden_states = hidden_states.chunk(2, dim=-1)
+ hidden_states = self.activation(chunked_hidden_states[0]) * chunked_hidden_states[1]
+ hidden_states = self.output_linear(hidden_states)
+ return hidden_states
+
+
+class GraniteMoeSharedDecoderLayer(GraniteMoeDecoderLayer):
+ def __init__(self, config: GraniteMoeSharedConfig, layer_idx: int):
+ super().__init__(config, layer_idx)
+ self.shared_mlp = None if config.shared_intermediate_size == 0 else GraniteMoeSharedMLP(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ output_attentions: bool | None = False,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[GraniteFlashAttentionKwargs],
+ ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ output_attentions=output_attentions,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+
+ hidden_states = residual + hidden_states * self.residual_multiplier
+
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ moe_hidden_states = self.block_sparse_moe(hidden_states)
+
+ if self.shared_mlp is None:
+ hidden_states = moe_hidden_states
+ else:
+ hidden_states = moe_hidden_states + self.shared_mlp(hidden_states)
+ hidden_states = residual + hidden_states * self.residual_multiplier
+ return hidden_states
+
+
+class GraniteMoeSharedPreTrainedModel(GraniteMoePreTrainedModel):
+ config: GraniteMoeSharedConfig
+ _no_split_modules = ["GraniteMoeSharedDecoderLayer"]
+
+
+class GraniteMoeSharedModel(GraniteMoeModel):
+ def __init__(self, config: GraniteMoeSharedConfig):
+ super().__init__(config)
+ self.layers = nn.ModuleList(
+ [GraniteMoeSharedDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+
+
+class GraniteMoeSharedForCausalLM(GraniteMoeForCausalLM):
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+ def __init__(self, config: GraniteMoeSharedConfig):
+ super().__init__(config)
+ self.model = GraniteMoeSharedModel(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+
+__all__ = ["GraniteMoeSharedForCausalLM", "GraniteMoeSharedModel", "GraniteMoeSharedPreTrainedModel"]
diff --git a/third_party/transformers/src/transformers/models/groupvit/__init__.py b/third_party/transformers/src/transformers/models/groupvit/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..10c315e280152759a30b035916cb7f43d6d40d16
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/groupvit/__init__.py
@@ -0,0 +1,27 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_groupvit import *
+ from .modeling_groupvit import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/groupvit/configuration_groupvit.py b/third_party/transformers/src/transformers/models/groupvit/configuration_groupvit.py
new file mode 100644
index 0000000000000000000000000000000000000000..9f02eb7e1e4f790ef30905490f00478a3a941584
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/groupvit/configuration_groupvit.py
@@ -0,0 +1,229 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""GroupViT model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="nvidia/groupvit-gcc-yfcc")
+@strict
+class GroupViTTextConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import GroupViTTextConfig, GroupViTTextModel
+
+ >>> # Initializing a GroupViTTextModel with nvidia/groupvit-gcc-yfcc style configuration
+ >>> configuration = GroupViTTextConfig()
+
+ >>> model = GroupViTTextModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "groupvit_text_model"
+ base_config_key = "text_config"
+
+ vocab_size: int = 49408
+ hidden_size: int = 256
+ intermediate_size: int = 1024
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 4
+ max_position_embeddings: int = 77
+ hidden_act: str = "quick_gelu"
+ layer_norm_eps: float = 1e-5
+ dropout: float | int = 0.0
+ attention_dropout: float | int = 0.0
+ initializer_range: float = 0.02
+ initializer_factor: float = 1.0
+ pad_token_id: int | None = 1
+ bos_token_id: int | None = 49406
+ eos_token_id: int | list[int] | None = 49407
+
+
+@auto_docstring(checkpoint="nvidia/groupvit-gcc-yfcc")
+@strict
+class GroupViTVisionConfig(PreTrainedConfig):
+ r"""
+ depths (`list[int]`, *optional*, defaults to [6, 3, 3]):
+ The number of layers in each encoder block.
+ num_group_tokens (`list[int]`, *optional*, defaults to [64, 8, 0]):
+ The number of group tokens for each stage.
+ num_output_groups (`list[int]`, *optional*, defaults to [64, 8, 8]):
+ The number of output groups for each stage, 0 means no group.
+ assign_eps (`float`, *optional*, defaults to `1.0`):
+ Epsilon used in layer norm
+ assign_mlp_ratio (`list[int]`, *optional*, defaults to `[0.5, 4]`):
+ Ratio used to infer hidden size of MLP layers.
+
+ Example:
+
+ ```python
+ >>> from transformers import GroupViTVisionConfig, GroupViTVisionModel
+
+ >>> # Initializing a GroupViTVisionModel with nvidia/groupvit-gcc-yfcc style configuration
+ >>> configuration = GroupViTVisionConfig()
+
+ >>> model = GroupViTVisionModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "groupvit_vision_model"
+ base_config_key = "vision_config"
+
+ hidden_size: int = 384
+ intermediate_size: int = 1536
+ num_hidden_layers: int = 12
+ depths: list[int] | tuple[int, ...] = (6, 3, 3)
+ num_group_tokens: list[int] | tuple[int, ...] = (64, 8, 0)
+ num_output_groups: list[int] | tuple[int, ...] = (64, 8, 8)
+ num_attention_heads: int = 6
+ image_size: int | list[int] | tuple[int, int] = 224
+ patch_size: int | list[int] | tuple[int, int] = 16
+ num_channels: int = 3
+ hidden_act: str = "gelu"
+ layer_norm_eps: float = 1e-5
+ dropout: float | int = 0.0
+ attention_dropout: float | int = 0.0
+ initializer_range: float = 0.02
+ initializer_factor: float = 1.0
+ assign_eps: float = 1.0
+ assign_mlp_ratio: list[float | int] | tuple[float | int, ...] = (0.5, 4)
+
+ def validate_architecture(self):
+ """Part of `@strict`-powered validation. Validates the architecture of the config."""
+ if self.num_hidden_layers != sum(self.depths):
+ logger.warning(
+ f"Manually setting num_hidden_layers to {self.num_hidden_layers}, but we expect num_hidden_layers ="
+ f" sum(depth) = {sum(self.depths)}"
+ )
+
+
+@auto_docstring(checkpoint="nvidia/groupvit-gcc-yfcc")
+@strict
+class GroupViTConfig(PreTrainedConfig):
+ r"""
+ projection_intermediate_dim (`int`, *optional*, defaults to 4096):
+ Dimensionality of intermediate layer of text and vision projection layers.
+ output_segmentation (`bool`, *optional*, defaults to False):
+ Whether or not to return the segmentation logits.
+ """
+
+ model_type = "groupvit"
+ sub_configs = {"text_config": GroupViTTextConfig, "vision_config": GroupViTVisionConfig}
+
+ text_config: dict | PreTrainedConfig | None = None
+ vision_config: dict | PreTrainedConfig | None = None
+ projection_dim: int = 256
+ projection_intermediate_dim: int = 4096
+ logit_scale_init_value: float = 2.6592
+ initializer_range: float = 0.02
+ initializer_factor: float = 1.0
+ output_segmentation: bool = False
+
+ def __post_init__(self, **kwargs):
+ if self.text_config is None:
+ text_config = {}
+ logger.info("`text_config` is `None`. Initializing the `GroupViTTextConfig` with default values.")
+ elif isinstance(self.text_config, GroupViTTextConfig):
+ text_config = self.text_config.to_dict()
+ else:
+ text_config = self.text_config
+
+ if self.vision_config is None:
+ vision_config = {}
+ logger.info("`vision_config` is `None`. initializing the `GroupViTVisionConfig` with default values.")
+ elif isinstance(self.vision_config, GroupViTVisionConfig):
+ vision_config = self.vision_config.to_dict()
+ else:
+ vision_config = self.vision_config
+
+ # For backward compatibility check keyword args
+ # Instead of simply assigning `[text|vision]_config_dict` to `[text|vision]_config`, we use the values in
+ # `[text|vision]_config_dict` to update the values in `[text|vision]_config`. The values should be same in most
+ # cases, but we don't want to break anything regarding `_config_dict` that existed before commit `8827e1b2`.
+ text_config_dict = kwargs.pop("text_config_dict", None)
+ vision_config_dict = kwargs.pop("vision_config_dict", None)
+
+ if text_config_dict is not None:
+ # This is the complete result when using `text_config_dict`.
+ _text_config_dict = GroupViTTextConfig(**text_config_dict).to_dict()
+
+ # Give a warning if the values exist in both `_text_config_dict` and `text_config` but being different.
+ for key, value in _text_config_dict.items():
+ if key in text_config and value != text_config[key] and key != "transformers_version":
+ # If specified in `text_config_dict`
+ if key in text_config_dict:
+ message = (
+ f"`{key}` is found in both `text_config_dict` and `text_config` but with different values. "
+ f'The value `text_config_dict["{key}"]` will be used instead.'
+ )
+ # If inferred from default argument values (just to be super careful)
+ else:
+ message = (
+ f"`text_config_dict` is provided which will be used to initialize `GroupViTTextConfig`. The "
+ f'value `text_config["{key}"]` will be overridden.'
+ )
+ logger.info(message)
+
+ # Update all values in `text_config` with the ones in `_text_config_dict`.
+ text_config.update(_text_config_dict)
+
+ if vision_config_dict is not None:
+ # This is the complete result when using `vision_config_dict`.
+ _vision_config_dict = GroupViTVisionConfig(**vision_config_dict).to_dict()
+ # convert keys to string instead of integer
+ if "id2label" in _vision_config_dict:
+ _vision_config_dict["id2label"] = {
+ str(key): value for key, value in _vision_config_dict["id2label"].items()
+ }
+
+ # Give a warning if the values exist in both `_vision_config_dict` and `vision_config` but being different.
+ for key, value in _vision_config_dict.items():
+ if key in vision_config and value != vision_config[key] and key != "transformers_version":
+ # If specified in `vision_config_dict`
+ if key in vision_config_dict:
+ message = (
+ f"`{key}` is found in both `vision_config_dict` and `vision_config` but with different "
+ f'values. The value `vision_config_dict["{key}"]` will be used instead.'
+ )
+ # If inferred from default argument values (just to be super careful)
+ else:
+ message = (
+ f"`vision_config_dict` is provided which will be used to initialize `GroupViTVisionConfig`. "
+ f'The value `vision_config["{key}"]` will be overridden.'
+ )
+ logger.info(message)
+
+ # Update all values in `vision_config` with the ones in `_vision_config_dict`.
+ vision_config.update(_vision_config_dict)
+
+ # Finally we can convert back our unified text/vision configs to `PretrainedConfig`
+ self.text_config = GroupViTTextConfig(**text_config)
+ self.vision_config = GroupViTVisionConfig(**vision_config)
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["GroupViTConfig", "GroupViTTextConfig", "GroupViTVisionConfig"]
diff --git a/third_party/transformers/src/transformers/models/groupvit/convert_groupvit_nvlab_to_hf.py b/third_party/transformers/src/transformers/models/groupvit/convert_groupvit_nvlab_to_hf.py
new file mode 100644
index 0000000000000000000000000000000000000000..57d2759d65fa2e6c8e35631dbc60fdf162818a64
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/groupvit/convert_groupvit_nvlab_to_hf.py
@@ -0,0 +1,218 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+
+"""
+Convert GroupViT checkpoints from the original repository.
+
+URL: https://github.com/NVlabs/GroupViT
+"""
+
+import argparse
+from io import BytesIO
+
+import httpx
+import torch
+from PIL import Image
+
+from transformers import CLIPProcessor, GroupViTConfig, GroupViTModel
+
+
+def rename_key(name):
+ # vision encoder
+ if "img_encoder.pos_embed" in name:
+ name = name.replace("img_encoder.pos_embed", "vision_model.embeddings.position_embeddings")
+ if "img_encoder.patch_embed.proj" in name:
+ name = name.replace("img_encoder.patch_embed.proj", "vision_model.embeddings.patch_embeddings.projection")
+ if "img_encoder.patch_embed.norm" in name:
+ name = name.replace("img_encoder.patch_embed.norm", "vision_model.embeddings.layernorm")
+ if "img_encoder.layers" in name:
+ name = name.replace("img_encoder.layers", "vision_model.encoder.stages")
+ if "blocks" in name and "res" not in name:
+ name = name.replace("blocks", "layers")
+ if "attn" in name and "pre_assign" not in name:
+ name = name.replace("attn", "self_attn")
+ if "proj" in name and "self_attn" in name and "text" not in name:
+ name = name.replace("proj", "out_proj")
+ if "pre_assign_attn.attn.proj" in name:
+ name = name.replace("pre_assign_attn.attn.proj", "pre_assign_attn.attn.out_proj")
+ if "norm1" in name:
+ name = name.replace("norm1", "layer_norm1")
+ if "norm2" in name and "pre_assign" not in name:
+ name = name.replace("norm2", "layer_norm2")
+ if "img_encoder.norm" in name:
+ name = name.replace("img_encoder.norm", "vision_model.layernorm")
+ # text encoder
+ if "text_encoder.token_embedding" in name:
+ name = name.replace("text_encoder.token_embedding", "text_model.embeddings.token_embedding")
+ if "text_encoder.positional_embedding" in name:
+ name = name.replace("text_encoder.positional_embedding", "text_model.embeddings.position_embedding.weight")
+ if "text_encoder.transformer.resblocks." in name:
+ name = name.replace("text_encoder.transformer.resblocks.", "text_model.encoder.layers.")
+ if "ln_1" in name:
+ name = name.replace("ln_1", "layer_norm1")
+ if "ln_2" in name:
+ name = name.replace("ln_2", "layer_norm2")
+ if "c_fc" in name:
+ name = name.replace("c_fc", "fc1")
+ if "c_proj" in name:
+ name = name.replace("c_proj", "fc2")
+ if "text_encoder" in name:
+ name = name.replace("text_encoder", "text_model")
+ if "ln_final" in name:
+ name = name.replace("ln_final", "final_layer_norm")
+ # projection layers
+ if "img_projector.linear_hidden." in name:
+ name = name.replace("img_projector.linear_hidden.", "visual_projection.")
+ if "img_projector.linear_out." in name:
+ name = name.replace("img_projector.linear_out.", "visual_projection.3.")
+ if "text_projector.linear_hidden" in name:
+ name = name.replace("text_projector.linear_hidden", "text_projection")
+ if "text_projector.linear_out" in name:
+ name = name.replace("text_projector.linear_out", "text_projection.3")
+
+ return name
+
+
+def convert_state_dict(orig_state_dict, config):
+ for key in orig_state_dict.copy():
+ val = orig_state_dict.pop(key)
+
+ if "qkv" in key:
+ # weights and biases of the key, value and query projections of vision encoder's attention layers require special treatment:
+ # we need to split them up into separate matrices/vectors
+ key_split = key.split(".")
+ stage_num, layer_num = int(key_split[2]), int(key_split[4])
+ dim = config.vision_config.hidden_size
+ if "weight" in key:
+ orig_state_dict[
+ f"vision_model.encoder.stages.{stage_num}.layers.{layer_num}.self_attn.q_proj.weight"
+ ] = val[:dim, :]
+ orig_state_dict[
+ f"vision_model.encoder.stages.{stage_num}.layers.{layer_num}.self_attn.k_proj.weight"
+ ] = val[dim : dim * 2, :]
+ orig_state_dict[
+ f"vision_model.encoder.stages.{stage_num}.layers.{layer_num}.self_attn.v_proj.weight"
+ ] = val[-dim:, :]
+ else:
+ orig_state_dict[
+ f"vision_model.encoder.stages.{stage_num}.layers.{layer_num}.self_attn.q_proj.bias"
+ ] = val[:dim]
+ orig_state_dict[
+ f"vision_model.encoder.stages.{stage_num}.layers.{layer_num}.self_attn.k_proj.bias"
+ ] = val[dim : dim * 2]
+ orig_state_dict[
+ f"vision_model.encoder.stages.{stage_num}.layers.{layer_num}.self_attn.v_proj.bias"
+ ] = val[-dim:]
+ elif "in_proj" in key:
+ # weights and biases of the key, value and query projections of text encoder's attention layers require special treatment:
+ # we need to split them up into separate matrices/vectors
+ key_split = key.split(".")
+ layer_num = int(key_split[3])
+ dim = config.text_config.hidden_size
+ if "weight" in key:
+ orig_state_dict[f"text_model.encoder.layers.{layer_num}.self_attn.q_proj.weight"] = val[:dim, :]
+ orig_state_dict[f"text_model.encoder.layers.{layer_num}.self_attn.k_proj.weight"] = val[
+ dim : dim * 2, :
+ ]
+ orig_state_dict[f"text_model.encoder.layers.{layer_num}.self_attn.v_proj.weight"] = val[-dim:, :]
+ else:
+ orig_state_dict[f"text_model.encoder.layers.{layer_num}.self_attn.q_proj.bias"] = val[:dim]
+ orig_state_dict[f"text_model.encoder.layers.{layer_num}.self_attn.k_proj.bias"] = val[dim : dim * 2]
+ orig_state_dict[f"text_model.encoder.layers.{layer_num}.self_attn.v_proj.bias"] = val[-dim:]
+ else:
+ new_name = rename_key(key)
+ # squeeze if necessary
+ if (
+ "text_projection.0" in new_name
+ or "text_projection.3" in new_name
+ or "visual_projection.0" in new_name
+ or "visual_projection.3" in new_name
+ ):
+ orig_state_dict[new_name] = val.squeeze_()
+ else:
+ orig_state_dict[new_name] = val
+
+ return orig_state_dict
+
+
+# We will verify our results on an image of cute cats
+def prepare_img():
+ url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ with httpx.stream("GET", url) as response:
+ image = Image.open(BytesIO(response.read()))
+ return image
+
+
+@torch.no_grad()
+def convert_groupvit_checkpoint(
+ checkpoint_path, pytorch_dump_folder_path, model_name="groupvit-gcc-yfcc", push_to_hub=False
+):
+ """
+ Copy/paste/tweak model's weights to the Transformers design.
+ """
+ config = GroupViTConfig()
+ model = GroupViTModel(config).eval()
+
+ state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)["model"]
+ new_state_dict = convert_state_dict(state_dict, config)
+ missing_keys, unexpected_keys = model.load_state_dict(new_state_dict, strict=False)
+ assert missing_keys == ["text_model.embeddings.position_ids"]
+ assert (unexpected_keys == ["multi_label_logit_scale"]) or (len(unexpected_keys) == 0)
+
+ # verify result
+ processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+ image = prepare_img()
+ inputs = processor(text=["a photo of a cat", "a photo of a dog"], images=image, padding=True, return_tensors="pt")
+
+ with torch.no_grad():
+ outputs = model(**inputs)
+
+ if model_name == "groupvit-gcc-yfcc":
+ expected_logits = torch.tensor([[13.3523, 6.3629]])
+ elif model_name == "groupvit-gcc-redcaps":
+ expected_logits = torch.tensor([[16.1873, 8.6230]])
+ else:
+ raise ValueError(f"Model name {model_name} not supported.")
+ assert torch.allclose(outputs.logits_per_image, expected_logits, atol=1e-3)
+
+ processor.save_pretrained(pytorch_dump_folder_path)
+ model.save_pretrained(pytorch_dump_folder_path)
+ print("Successfully saved processor and model to", pytorch_dump_folder_path)
+
+ if push_to_hub:
+ print("Pushing to the hub...")
+ processor.push_to_hub(repo_id=f"nielsr/{model_name}")
+ model.push_to_hub(repo_id=f"nielsr/{model_name}")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument(
+ "--pytorch_dump_folder_path", default=None, type=str, help="Path to dump the processor and PyTorch model."
+ )
+ parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to GroupViT checkpoint")
+ parser.add_argument(
+ "--model_name",
+ default="groupvit-gccy-fcc",
+ type=str,
+ help="Name of the model. Expecting either 'groupvit-gcc-yfcc' or 'groupvit-gcc-redcaps'",
+ )
+ parser.add_argument(
+ "--push_to_hub",
+ action="store_true",
+ help="Whether or not to push the converted model and processor to the Hugging Face hub using the provided `model_name`.",
+ )
+ args = parser.parse_args()
+
+ convert_groupvit_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.model_name, args.push_to_hub)
diff --git a/third_party/transformers/src/transformers/models/groupvit/modeling_groupvit.py b/third_party/transformers/src/transformers/models/groupvit/modeling_groupvit.py
new file mode 100644
index 0000000000000000000000000000000000000000..b242c485df4c9fc8ed864b8812153144000c171a
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/groupvit/modeling_groupvit.py
@@ -0,0 +1,1343 @@
+# Copyright 2022 NVIDIA and The HuggingFace Team. All rights reserved.
+#
+# 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 GroupViT model."""
+
+import collections.abc
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...masking_utils import create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
+from ...modeling_utils import PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_groupvit import GroupViTConfig, GroupViTTextConfig, GroupViTVisionConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# contrastive loss function, adapted from
+# https://sachinruk.github.io/blog/pytorch/pytorch%20lightning/loss%20function/gpu/2021/03/07/CLIP.html
+def contrastive_loss(logits: torch.Tensor) -> torch.Tensor:
+ return nn.functional.cross_entropy(logits, torch.arange(len(logits), device=logits.device))
+
+
+# Copied from transformers.models.clip.modeling_clip.clip_loss with clip->groupvit
+def groupvit_loss(similarity: torch.Tensor) -> torch.Tensor:
+ caption_loss = contrastive_loss(similarity)
+ image_loss = contrastive_loss(similarity.t())
+ return (caption_loss + image_loss) / 2.0
+
+
+def hard_softmax(logits: torch.Tensor, dim: int):
+ y_soft = logits.softmax(dim)
+ # Straight through.
+ index = y_soft.max(dim, keepdim=True)[1]
+ y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
+ ret = y_hard - y_soft.detach() + y_soft
+
+ return ret
+
+
+def gumbel_softmax(logits: torch.Tensor, tau: float = 1, hard: bool = False, dim: int = -1) -> torch.Tensor:
+ # more stable https://github.com/pytorch/pytorch/issues/41663
+ gumbel_dist = torch.distributions.gumbel.Gumbel(
+ torch.tensor(0.0, device=logits.device, dtype=logits.dtype),
+ torch.tensor(1.0, device=logits.device, dtype=logits.dtype),
+ )
+ gumbels = gumbel_dist.sample(logits.shape)
+
+ gumbels = (logits + gumbels) / tau # ~Gumbel(logits,tau)
+ y_soft = gumbels.softmax(dim)
+
+ if hard:
+ # Straight through.
+ index = y_soft.max(dim, keepdim=True)[1]
+ y_hard = torch.zeros_like(logits, memory_format=torch.legacy_contiguous_format).scatter_(dim, index, 1.0)
+ ret = y_hard - y_soft.detach() + y_soft
+ else:
+ # Reparameterization trick.
+ ret = y_soft
+ return ret
+
+
+def resize_attention_map(attentions, height, width, align_corners=False):
+ """
+ Args:
+ attentions (`torch.Tensor`): attention map of shape [batch_size, groups, feat_height*feat_width]
+ height (`int`): height of the output attention map
+ width (`int`): width of the output attention map
+ align_corners (`bool`, *optional*): the `align_corner` argument for `nn.functional.interpolate`.
+
+ Returns:
+ `torch.Tensor`: resized attention map of shape [batch_size, groups, height, width]
+ """
+
+ scale = (height * width // attentions.shape[2]) ** 0.5
+ if height > width:
+ feat_width = int(np.round(width / scale))
+ feat_height = attentions.shape[2] // feat_width
+ else:
+ feat_height = int(np.round(height / scale))
+ feat_width = attentions.shape[2] // feat_height
+
+ batch_size = attentions.shape[0]
+ groups = attentions.shape[1] # number of group token
+ # [batch_size, groups, height*width, groups] -> [batch_size, groups, height, width]
+ attentions = attentions.reshape(batch_size, groups, feat_height, feat_width)
+ attentions = nn.functional.interpolate(
+ attentions, size=(height, width), mode="bilinear", align_corners=align_corners
+ )
+ return attentions
+
+
+def get_grouping_from_attentions(attentions, hw_shape):
+ """
+ Args:
+ attentions (`tuple(torch.FloatTensor)`: tuple of attention maps returned by `GroupViTVisionTransformer`
+ hw_shape (`tuple(int)`): height and width of the output attention map
+ Returns:
+ `torch.Tensor`: the attention map of shape [batch_size, groups, height, width]
+ """
+
+ attn_maps = []
+ with torch.no_grad():
+ prev_attn_masks = None
+ for attn_masks in attentions:
+ # [batch_size, num_groups, height x width] -> [batch_size, height x width, num_groups]
+ attn_masks = attn_masks.permute(0, 2, 1).contiguous()
+ if prev_attn_masks is None:
+ prev_attn_masks = attn_masks
+ else:
+ prev_attn_masks = prev_attn_masks @ attn_masks
+ # [batch_size, heightxwidth, num_groups] -> [batch_size, num_groups, heightxwidth] -> [batch_size, num_groups, height, width]
+ cur_attn_map = resize_attention_map(prev_attn_masks.permute(0, 2, 1).contiguous(), *hw_shape)
+ attn_maps.append(cur_attn_map)
+
+ # [batch_size, num_groups, height, width]
+ final_grouping = attn_maps[-1]
+
+ return final_grouping
+
+
+class GroupViTCrossAttentionLayer(nn.Module):
+ def __init__(self, config: GroupViTVisionConfig):
+ super().__init__()
+ self.attn = GroupViTAttention(config)
+ self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.mlp = GroupViTMLP(config)
+ self.norm_post = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+ def forward(self, query, key):
+ x = query
+ x = x + self.attn(query, encoder_hidden_states=key)[0]
+ x = x + self.mlp(self.norm2(x))
+ x = self.norm_post(x)
+ return x
+
+
+class GroupViTAssignAttention(nn.Module):
+ def __init__(self, config: GroupViTVisionConfig):
+ super().__init__()
+ self.scale = config.hidden_size**-0.5
+
+ self.q_proj = nn.Linear(config.hidden_size, config.hidden_size)
+ self.k_proj = nn.Linear(config.hidden_size, config.hidden_size)
+ self.v_proj = nn.Linear(config.hidden_size, config.hidden_size)
+ self.proj = nn.Linear(config.hidden_size, config.hidden_size)
+ self.assign_eps = config.assign_eps
+
+ def get_attn(self, attn, gumbel=True, hard=True):
+ if gumbel and self.training:
+ attn = gumbel_softmax(attn, dim=-2, hard=hard)
+ else:
+ if hard:
+ attn = hard_softmax(attn, dim=-2)
+ else:
+ attn = nn.functional.softmax(attn, dim=-2)
+
+ return attn
+
+ def forward(self, query, key):
+ value = key
+ # [batch_size, query_length, channels]
+ query = self.q_proj(query)
+
+ # [batch_size, key_length, channels]
+ key = self.k_proj(key)
+
+ # [batch_size, key_length, channels]
+ value = self.v_proj(value)
+
+ # [batch_size, query_length, key_length]
+ raw_attn = (query @ key.transpose(-2, -1)) * self.scale
+
+ attn = self.get_attn(raw_attn)
+ soft_attn = self.get_attn(raw_attn, gumbel=False, hard=False)
+
+ attn = attn / (attn.sum(dim=-1, keepdim=True) + self.assign_eps)
+
+ out = attn @ value
+
+ out = self.proj(out)
+
+ return out, soft_attn
+
+
+class GroupViTTokenAssign(nn.Module):
+ def __init__(self, config: GroupViTVisionConfig, num_group_token, num_output_group):
+ super().__init__()
+ self.num_output_group = num_output_group
+ # norm on group_tokens
+ self.norm_tokens = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ assign_mlp_ratio = (
+ config.assign_mlp_ratio
+ if isinstance(config.assign_mlp_ratio, collections.abc.Iterable)
+ else (config.assign_mlp_ratio, config.assign_mlp_ratio)
+ )
+ tokens_dim, channels_dim = [int(x * config.hidden_size) for x in assign_mlp_ratio]
+ self.mlp_inter = GroupViTMixerMLP(config, num_group_token, tokens_dim, num_output_group)
+ self.norm_post_tokens = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ # norm on x
+ self.norm_x = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.pre_assign_attn = GroupViTCrossAttentionLayer(config)
+
+ self.assign = GroupViTAssignAttention(config)
+ self.norm_new_x = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.mlp_channels = GroupViTMLP(config, config.hidden_size, channels_dim, config.hidden_size)
+
+ def project_group_token(self, group_tokens):
+ """
+ Args:
+ group_tokens (torch.Tensor): group tokens, [batch_size, num_group_tokens, channels]
+
+ Returns:
+ projected_group_tokens (torch.Tensor): [batch_size, num_output_groups, channels]
+ """
+ # [B, num_output_groups, C] <- [B, num_group_tokens, C]
+ projected_group_tokens = self.mlp_inter(group_tokens)
+ projected_group_tokens = self.norm_post_tokens(projected_group_tokens)
+ return projected_group_tokens
+
+ def forward(self, image_tokens, group_tokens):
+ """
+ Args:
+ image_tokens (`torch.Tensor`): image tokens, of shape [batch_size, input_length, channels]
+ group_tokens (`torch.Tensor`): group tokens, [batch_size, num_group_tokens, channels]
+ """
+
+ group_tokens = self.norm_tokens(group_tokens)
+ image_tokens = self.norm_x(image_tokens)
+ # [batch_size, num_output_groups, channels]
+ projected_group_tokens = self.project_group_token(group_tokens)
+ projected_group_tokens = self.pre_assign_attn(projected_group_tokens, image_tokens)
+ new_image_tokens, attention = self.assign(projected_group_tokens, image_tokens)
+ new_image_tokens += projected_group_tokens
+
+ new_image_tokens = new_image_tokens + self.mlp_channels(self.norm_new_x(new_image_tokens))
+
+ return new_image_tokens, attention
+
+
+@dataclass
+@auto_docstring
+class GroupViTModelOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `return_loss` is `True`):
+ Contrastive loss for image-text similarity.
+ logits_per_image (`torch.FloatTensor` of shape `(image_batch_size, text_batch_size)`):
+ The scaled dot product scores between `image_embeds` and `text_embeds`. This represents the image-text
+ similarity scores.
+ logits_per_text (`torch.FloatTensor` of shape `(text_batch_size, image_batch_size)`):
+ The scaled dot product scores between `text_embeds` and `image_embeds`. This represents the text-image
+ similarity scores.
+ segmentation_logits (`torch.FloatTensor` of shape `(batch_size, config.num_labels, logits_height, logits_width)`):
+ Classification scores for each pixel.
+
+
+
+ The logits returned do not necessarily have the same size as the `pixel_values` passed as inputs. This is
+ to avoid doing two interpolations and lose some quality when a user needs to resize the logits to the
+ original image size as post-processing. You should always check your logits shape and resize as needed.
+
+
+ text_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):
+ The text embeddings obtained by applying the projection layer to the pooled output of
+ [`GroupViTTextModel`].
+ image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim`):
+ The image embeddings obtained by applying the projection layer to the pooled output of
+ [`GroupViTVisionModel`].
+ text_model_output (`BaseModelOutputWithPooling`):
+ The output of the [`GroupViTTextModel`].
+ vision_model_output (`BaseModelOutputWithPooling`):
+ The output of the [`GroupViTVisionModel`].
+ """
+
+ loss: torch.FloatTensor | None = None
+ logits_per_image: torch.FloatTensor | None = None
+ logits_per_text: torch.FloatTensor | None = None
+ segmentation_logits: torch.FloatTensor | None = None
+ text_embeds: torch.FloatTensor | None = None
+ image_embeds: torch.FloatTensor | None = None
+ text_model_output: BaseModelOutputWithPooling = None
+ vision_model_output: BaseModelOutputWithPooling = None
+
+ def to_tuple(self) -> tuple[Any]:
+ return tuple(
+ self[k] if k not in ["text_model_output", "vision_model_output"] else getattr(self, k).to_tuple()
+ for k in self.keys()
+ )
+
+
+class GroupViTPatchEmbeddings(nn.Module):
+ """
+ Image to Patch Embedding.
+ """
+
+ def __init__(
+ self,
+ image_size: int | list[int] | tuple[int, int] = 224,
+ patch_size: int | tuple[int, int] = 16,
+ num_channels: int = 3,
+ embed_dim: int = 768,
+ ):
+ super().__init__()
+ image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
+ patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
+ num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
+ self.image_size = image_size
+ self.patch_size = patch_size
+ self.num_patches = num_patches
+
+ self.projection = nn.Conv2d(num_channels, embed_dim, kernel_size=patch_size, stride=patch_size)
+
+ def forward(self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False) -> torch.Tensor:
+ batch_size, num_channels, height, width = pixel_values.shape
+ if not interpolate_pos_encoding:
+ if height != self.image_size[0] or width != self.image_size[1]:
+ raise ValueError(
+ f"Input image size ({height}*{width}) doesn't match model"
+ f" ({self.image_size[0]}*{self.image_size[1]})."
+ )
+ x = self.projection(pixel_values).flatten(2).transpose(1, 2)
+ return x
+
+
+class GroupViTVisionEmbeddings(nn.Module):
+ def __init__(self, config: GroupViTVisionConfig):
+ super().__init__()
+
+ self.patch_embeddings = GroupViTPatchEmbeddings(
+ image_size=config.image_size,
+ patch_size=config.patch_size,
+ num_channels=config.num_channels,
+ embed_dim=config.hidden_size,
+ )
+ num_patches = self.patch_embeddings.num_patches
+ self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches, config.hidden_size))
+ self.dropout = nn.Dropout(config.dropout)
+ self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.patch_size = config.patch_size
+ self.config = config
+
+ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+ """
+ This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution
+ images. This method is also adapted to support torch.jit tracing and no class embeddings.
+
+ Adapted from:
+ - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
+ - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
+ """
+
+ num_patches = embeddings.shape[1]
+ num_positions = self.position_embeddings.shape[1]
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embeddings
+
+ patch_pos_embed = self.position_embeddings
+
+ dim = embeddings.shape[-1]
+
+ new_height = height // self.patch_size
+ new_width = width // self.patch_size
+
+ sqrt_num_positions = torch_int(num_positions**0.5)
+ patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
+ patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed,
+ size=(new_height, new_width),
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+ return patch_pos_embed
+
+ def forward(self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False) -> torch.Tensor:
+ batch_size, num_channels, height, width = pixel_values.shape
+ embeddings = self.patch_embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
+
+ embeddings = self.layernorm(embeddings)
+
+ batch_size, seq_len, _ = embeddings.size()
+
+ # add positional encoding to each token
+ if interpolate_pos_encoding:
+ embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+ else:
+ embeddings = embeddings + self.position_embeddings
+
+ embeddings = self.dropout(embeddings)
+
+ return embeddings
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPTextEmbeddings with CLIP->GroupViT
+class GroupViTTextEmbeddings(nn.Module):
+ def __init__(self, config: GroupViTTextConfig):
+ super().__init__()
+ embed_dim = config.hidden_size
+
+ self.token_embedding = nn.Embedding(config.vocab_size, embed_dim)
+ self.position_embedding = nn.Embedding(config.max_position_embeddings, embed_dim)
+
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+ self.register_buffer(
+ "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+ )
+
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ ) -> torch.Tensor:
+ seq_length = input_ids.shape[-1] if input_ids is not None else inputs_embeds.shape[-2]
+ max_position_embedding = self.position_embedding.weight.shape[0]
+
+ if seq_length > max_position_embedding:
+ raise ValueError(
+ f"Sequence length must be less than max_position_embeddings (got `sequence length`: "
+ f"{seq_length} and max_position_embeddings: {max_position_embedding}"
+ )
+
+ if position_ids is None:
+ position_ids = self.position_ids[:, :seq_length]
+
+ if inputs_embeds is None:
+ inputs_embeds = self.token_embedding(input_ids)
+
+ position_embeddings = self.position_embedding(position_ids)
+ embeddings = inputs_embeds + position_embeddings
+
+ return embeddings
+
+
+class GroupViTStage(nn.Module):
+ """This corresponds to the `GroupingLayer` class in the GroupViT implementation."""
+
+ def __init__(
+ self,
+ config: GroupViTVisionConfig,
+ depth: int,
+ num_prev_group_token: int,
+ num_group_token: int,
+ num_output_group: int,
+ ):
+ super().__init__()
+ self.depth = depth
+ self.num_group_token = num_group_token
+ if num_group_token > 0:
+ self.group_token = nn.Parameter(torch.zeros(1, num_group_token, config.hidden_size))
+ else:
+ self.group_token = None
+ self.layers = nn.ModuleList([GroupViTEncoderLayer(config) for _ in range(depth)])
+
+ if num_group_token > 0:
+ self.downsample = GroupViTTokenAssign(
+ config=config,
+ num_group_token=num_group_token,
+ num_output_group=num_output_group,
+ )
+ else:
+ self.downsample = None
+
+ if num_prev_group_token > 0 and num_group_token > 0:
+ self.group_projector = nn.Sequential(
+ nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps),
+ GroupViTMixerMLP(config, num_prev_group_token, config.hidden_size // 2, num_group_token),
+ )
+ else:
+ self.group_projector = None
+
+ @property
+ def with_group_token(self):
+ return self.group_token is not None
+
+ def split_x(self, x):
+ if self.with_group_token:
+ return x[:, : -self.num_group_token], x[:, -self.num_group_token :]
+ else:
+ return x, None
+
+ def concat_x(self, x: torch.Tensor, group_token: torch.Tensor | None = None) -> torch.Tensor:
+ if group_token is None:
+ return x
+ return torch.cat([x, group_token], dim=1)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ prev_group_token: torch.Tensor | None = None,
+ output_attentions: bool | None = False,
+ ) -> tuple[torch.FloatTensor]:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ `(config.encoder_attention_heads,)`.
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the grouping tensors of Grouping block.
+ """
+ if self.with_group_token:
+ group_token = self.group_token.expand(hidden_states.size(0), -1, -1)
+ if self.group_projector is not None:
+ group_token = group_token + self.group_projector(prev_group_token)
+ else:
+ group_token = None
+
+ x = hidden_states
+
+ cat_x = self.concat_x(x, group_token)
+ for layer in self.layers:
+ cat_x = layer(cat_x, attention_mask=None)
+
+ x, group_token = self.split_x(cat_x)
+
+ attention = None
+ if self.downsample is not None:
+ x, attention = self.downsample(x, group_token)
+
+ outputs = (x, group_token)
+ if output_attentions:
+ outputs = outputs + (attention,)
+
+ return outputs
+
+
+class GroupViTMLP(nn.Module):
+ def __init__(
+ self,
+ config: GroupViTVisionConfig,
+ hidden_size: int | None = None,
+ intermediate_size: int | None = None,
+ output_size: int | None = None,
+ ):
+ super().__init__()
+ self.config = config
+ self.activation_fn = ACT2FN[config.hidden_act]
+ hidden_size = hidden_size if hidden_size is not None else config.hidden_size
+ intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size
+ output_size = output_size if output_size is not None else hidden_size
+ self.fc1 = nn.Linear(hidden_size, intermediate_size)
+ self.fc2 = nn.Linear(intermediate_size, output_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ return hidden_states
+
+
+class GroupViTMixerMLP(GroupViTMLP):
+ def forward(self, x):
+ x = super().forward(x.transpose(1, 2))
+ return x.transpose(1, 2)
+
+
+class GroupViTAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = self.embed_dim // self.num_heads
+ if self.head_dim * self.num_heads != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+ f" {self.num_heads})."
+ )
+ self.scale = self.head_dim**-0.5
+ self.dropout = config.attention_dropout
+
+ self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+ 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,
+ attention_mask: torch.Tensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ **kwargs,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Input shape: Batch x Time x Channel"""
+
+ bsz, tgt_len, embed_dim = hidden_states.size()
+ is_cross_attention = encoder_hidden_states is not None
+
+ # get query proj
+ query_states = self.q_proj(hidden_states) * self.scale
+ if is_cross_attention:
+ key_states = self._shape(self.k_proj(encoder_hidden_states), -1, bsz)
+ value_states = self._shape(self.v_proj(encoder_hidden_states), -1, bsz)
+ else:
+ key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
+ value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
+
+ proj_shape = (bsz * self.num_heads, -1, self.head_dim)
+ query_states = self._shape(query_states, tgt_len, bsz).view(*proj_shape)
+ key_states = key_states.view(*proj_shape)
+ value_states = value_states.view(*proj_shape)
+
+ src_len = key_states.size(1)
+ attn_weights = torch.bmm(query_states, key_states.transpose(1, 2))
+
+ if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
+ raise ValueError(
+ f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is"
+ f" {attn_weights.size()}"
+ )
+
+ if attention_mask is not None:
+ if attention_mask.size() != (bsz, 1, tgt_len, src_len):
+ raise ValueError(
+ f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
+ )
+ attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
+ attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+
+ # this operation is a bit awkward, but it's required to
+ # make sure that attn_weights keeps its gradient.
+ # In order to do so, attn_weights have to reshaped
+ # twice and have to be reused in the following
+ attn_weights_reshaped = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+ attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len)
+
+ attn_probs = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)
+
+ attn_output = torch.bmm(attn_probs, value_states)
+
+ if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):
+ raise ValueError(
+ f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
+ f" {attn_output.size()}"
+ )
+
+ attn_output = attn_output.view(bsz, self.num_heads, tgt_len, self.head_dim)
+ attn_output = attn_output.transpose(1, 2)
+ attn_output = attn_output.reshape(bsz, tgt_len, embed_dim)
+
+ attn_output = self.out_proj(attn_output)
+
+ return attn_output, attn_weights_reshaped
+
+
+# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->GroupViT
+class GroupViTEncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: GroupViTConfig):
+ super().__init__()
+ self.embed_dim = config.hidden_size
+ self.self_attn = GroupViTAttention(config)
+ self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+ self.mlp = GroupViTMLP(config)
+ self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.FloatTensor, torch.Tensor | None]:
+ residual = hidden_states
+
+ hidden_states = self.layer_norm1(hidden_states)
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.layer_norm2(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+@auto_docstring
+class GroupViTPreTrainedModel(PreTrainedModel):
+ config: GroupViTConfig
+ base_model_prefix = "groupvit"
+ input_modalities = ("image", "text")
+ supports_gradient_checkpointing = True
+ _can_record_outputs = {
+ "hidden_states": GroupViTEncoderLayer,
+ "attentions": GroupViTAttention,
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ """Initialize the weights"""
+
+ init_range = self.config.initializer_range
+ if isinstance(module, (nn.Linear, nn.Conv2d)):
+ init.normal_(module.weight, mean=0.0, std=init_range)
+ if module.bias is not None:
+ init.zeros_(module.bias)
+ elif isinstance(module, (nn.LayerNorm, nn.BatchNorm1d)):
+ init.zeros_(module.bias)
+ init.ones_(module.weight)
+ if getattr(module, "running_mean", None) is not None:
+ init.zeros_(module.running_mean)
+ init.ones_(module.running_var)
+ init.zeros_(module.num_batches_tracked)
+
+ factor = self.config.initializer_factor
+ if isinstance(module, GroupViTTextEmbeddings):
+ init.normal_(module.token_embedding.weight, mean=0.0, std=factor * 0.02)
+ init.normal_(module.position_embedding.weight, mean=0.0, std=factor * 0.02)
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+ elif isinstance(module, GroupViTAttention):
+ factor = self.config.initializer_factor
+ in_proj_std = (module.embed_dim**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
+ out_proj_std = (module.embed_dim**-0.5) * factor
+ init.normal_(module.q_proj.weight, std=in_proj_std)
+ init.normal_(module.k_proj.weight, std=in_proj_std)
+ init.normal_(module.v_proj.weight, std=in_proj_std)
+ init.normal_(module.out_proj.weight, std=out_proj_std)
+ elif isinstance(module, GroupViTMLP):
+ factor = self.config.initializer_factor
+ in_proj_std = (module.config.hidden_size**-0.5) * ((2 * module.config.num_hidden_layers) ** -0.5) * factor
+ fc_std = (2 * module.config.hidden_size) ** -0.5 * factor
+ init.normal_(module.fc1.weight, std=fc_std)
+ init.normal_(module.fc2.weight, std=in_proj_std)
+
+
+class GroupViTVisionEncoder(nn.Module):
+ def __init__(self, config: GroupViTVisionConfig) -> None:
+ super().__init__()
+ self.config = config
+ self.stages = nn.ModuleList(
+ [
+ GroupViTStage(
+ config=config,
+ depth=config.depths[i],
+ num_group_token=config.num_group_tokens[i],
+ num_output_group=config.num_output_groups[i],
+ num_prev_group_token=config.num_output_groups[i - 1] if i > 0 else 0,
+ )
+ for i in range(len(config.depths))
+ ]
+ )
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ output_hidden_states: bool | None = None,
+ output_attentions: bool | None = None,
+ return_dict: bool | None = None,
+ ) -> tuple | BaseModelOutput:
+ 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.return_dict
+
+ all_hidden_states = () if output_hidden_states else None
+ all_groupings = () if output_attentions else None
+
+ group_tokens = None
+
+ for i, stage in enumerate(self.stages):
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ layer_outputs = stage(hidden_states, group_tokens, output_attentions)
+
+ hidden_states = layer_outputs[0]
+ group_tokens = layer_outputs[1]
+
+ if output_attentions and layer_outputs[2] is not None:
+ all_groupings = all_groupings + (layer_outputs[2],)
+
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ if not return_dict:
+ return tuple(v for v in [hidden_states, all_hidden_states, all_groupings] if v is not None)
+ return BaseModelOutput(
+ last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_groupings
+ )
+
+
+class GroupViTTextEncoder(nn.Module):
+ """
+ Transformer encoder consisting of `config.num_hidden_layers` self-attention layers. Each layer is a
+ [`GroupViTEncoderLayer`].
+
+ Args:
+ config: GroupViTTextConfig
+ """
+
+ def __init__(self, config: GroupViTTextConfig):
+ super().__init__()
+ self.config = config
+ self.layers = nn.ModuleList([GroupViTEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ inputs_embeds,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutput:
+ r"""
+ Args:
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ 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.
+ 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)
+ """
+ hidden_states = inputs_embeds
+ for encoder_layer in self.layers:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask,
+ **kwargs,
+ )
+
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+class GroupViTTextTransformer(GroupViTPreTrainedModel):
+ def __init__(self, config: GroupViTTextConfig):
+ super().__init__(config)
+ embed_dim = config.hidden_size
+ self.embeddings = GroupViTTextEmbeddings(config)
+ self.encoder = GroupViTTextEncoder(config)
+ self.final_layer_norm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+ # For `pooled_output` computation
+ self.eos_token_id = config.eos_token_id
+
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPooling:
+ if input_ids is None:
+ raise ValueError("You have to specify input_ids")
+
+ input_shape = input_ids.size()
+ input_ids = input_ids.view(-1, input_shape[-1])
+
+ hidden_states = self.embeddings(input_ids=input_ids, position_ids=position_ids)
+
+ attention_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=hidden_states,
+ attention_mask=attention_mask,
+ past_key_values=None,
+ )
+
+ kwargs.pop("is_causal", None)
+ encoder_outputs: BaseModelOutput = self.encoder(
+ inputs_embeds=hidden_states,
+ attention_mask=attention_mask,
+ is_causal=True,
+ **kwargs,
+ )
+
+ last_hidden_state = encoder_outputs[0]
+ last_hidden_state = self.final_layer_norm(last_hidden_state)
+
+ if self.eos_token_id == 2:
+ # The `eos_token_id` was incorrect before PR #24773: Let's keep what have been done here.
+ # A CLIP model with such `eos_token_id` in the config can't work correctly with extra new tokens added
+ # ------------------------------------------------------------
+ # text_embeds.shape = [batch_size, sequence_length, transformer.width]
+ # take features from the eot embedding (eot_token is the highest number in each sequence)
+ # casting to torch.int for onnx compatibility: argmax doesn't support int64 inputs with opset 14
+ pooled_output = last_hidden_state[
+ torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
+ input_ids.to(dtype=torch.int, device=last_hidden_state.device).argmax(dim=-1),
+ ]
+ else:
+ # The config gets updated `eos_token_id` from PR #24773 (so the use of extra new tokens is possible)
+ pooled_output = last_hidden_state[
+ torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device),
+ # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`)
+ # Note: we assume each sequence (along batch dim.) contains an `eos_token_id` (e.g. prepared by the tokenizer)
+ (input_ids.to(dtype=torch.int, device=last_hidden_state.device) == self.eos_token_id)
+ .int()
+ .argmax(dim=-1),
+ ]
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooled_output,
+ )
+
+
+class GroupViTTextModel(GroupViTPreTrainedModel):
+ config: GroupViTTextConfig
+ input_modalities = ("text",)
+
+ def __init__(self, config: GroupViTTextConfig):
+ super().__init__(config)
+ self.text_model = GroupViTTextTransformer(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Module:
+ return self.text_model.embeddings.token_embedding
+
+ def set_input_embeddings(self, value):
+ self.text_model.embeddings.token_embedding = value
+
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ Examples:
+
+ ```python
+ >>> from transformers import CLIPTokenizer, GroupViTTextModel
+
+ >>> tokenizer = CLIPTokenizer.from_pretrained("nvidia/groupvit-gcc-yfcc")
+ >>> model = GroupViTTextModel.from_pretrained("nvidia/groupvit-gcc-yfcc")
+
+ >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")
+
+ >>> outputs = model(**inputs)
+ >>> last_hidden_state = outputs.last_hidden_state
+ >>> pooled_output = outputs.pooler_output # pooled (EOS token) states
+ ```"""
+ return self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+
+class GroupViTVisionTransformer(nn.Module):
+ def __init__(self, config: GroupViTVisionConfig):
+ super().__init__()
+ self.config = config
+ embed_dim = config.hidden_size
+
+ self.embeddings = GroupViTVisionEmbeddings(config)
+ self.encoder = GroupViTVisionEncoder(config)
+ self.layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ output_hidden_states: bool | None = None,
+ output_attentions: bool | None = None,
+ return_dict: bool | None = None,
+ ) -> tuple | BaseModelOutputWithPooling:
+ 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.return_dict
+
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ hidden_states = self.embeddings(pixel_values)
+
+ encoder_outputs = self.encoder(
+ hidden_states=hidden_states,
+ output_hidden_states=output_hidden_states,
+ output_attentions=output_attentions,
+ return_dict=return_dict,
+ )
+
+ last_hidden_state = encoder_outputs[0]
+
+ # normalize the last hidden state
+ last_hidden_state = self.layernorm(last_hidden_state)
+ pooled_output = last_hidden_state.mean(dim=1)
+
+ if not return_dict:
+ return (last_hidden_state, pooled_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooled_output,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ )
+
+
+class GroupViTVisionModel(GroupViTPreTrainedModel):
+ config: GroupViTVisionConfig
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+ _can_record_outputs = {}
+
+ def __init__(self, config: GroupViTVisionConfig):
+ super().__init__(config)
+ self.vision_model = GroupViTVisionTransformer(config)
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self) -> GroupViTPatchEmbeddings:
+ return self.vision_model.embeddings.patch_embeddings
+
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ return_dict: bool | None = None,
+ **kwargs,
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ Examples:
+
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoProcessor, GroupViTVisionModel
+
+ >>> processor = AutoProcessor.from_pretrained("nvidia/groupvit-gcc-yfcc")
+ >>> model = GroupViTVisionModel.from_pretrained("nvidia/groupvit-gcc-yfcc")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> inputs = processor(images=image, return_tensors="pt")
+
+ >>> outputs = model(**inputs)
+ >>> last_hidden_state = outputs.last_hidden_state
+ >>> pooled_output = outputs.pooler_output # pooled CLS states
+ ```"""
+ return self.vision_model(
+ pixel_values=pixel_values,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+
+@auto_docstring
+class GroupViTModel(GroupViTPreTrainedModel):
+ config: GroupViTConfig
+
+ def __init__(self, config: GroupViTConfig):
+ super().__init__(config)
+
+ if not isinstance(config.text_config, GroupViTTextConfig):
+ raise TypeError(
+ "config.text_config is expected to be of type GroupViTTextConfig but is of type"
+ f" {type(config.text_config)}."
+ )
+
+ if not isinstance(config.vision_config, GroupViTVisionConfig):
+ raise TypeError(
+ "config.vision_config is expected to be of type GroupViTVisionConfig but is of type"
+ f" {type(config.vision_config)}."
+ )
+
+ text_config = config.text_config
+ vision_config = config.vision_config
+
+ self.projection_dim = config.projection_dim
+ self.projection_intermediate_dim = config.projection_intermediate_dim
+ self.text_embed_dim = text_config.hidden_size
+ self.vision_embed_dim = vision_config.hidden_size
+
+ self.text_model = GroupViTTextTransformer(text_config)
+ self.vision_model = GroupViTVisionTransformer(vision_config)
+
+ self.visual_projection = nn.Sequential(
+ nn.Linear(self.vision_embed_dim, self.projection_intermediate_dim, bias=True),
+ nn.BatchNorm1d(self.projection_intermediate_dim),
+ nn.ReLU(inplace=True),
+ nn.Linear(self.projection_intermediate_dim, self.projection_dim, bias=True),
+ )
+ self.text_projection = nn.Sequential(
+ nn.Linear(self.text_embed_dim, self.projection_intermediate_dim, bias=True),
+ nn.BatchNorm1d(self.projection_intermediate_dim),
+ nn.ReLU(inplace=True),
+ nn.Linear(self.projection_intermediate_dim, self.projection_dim, bias=True),
+ )
+ self.logit_scale = nn.Parameter(torch.tensor(self.config.logit_scale_init_value))
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def get_text_features(
+ self,
+ input_ids: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ Examples:
+
+ ```python
+ >>> import torch
+ >>> from transformers import CLIPTokenizer, GroupViTModel
+
+ >>> model = GroupViTModel.from_pretrained("nvidia/groupvit-gcc-yfcc")
+ >>> tokenizer = CLIPTokenizer.from_pretrained("nvidia/groupvit-gcc-yfcc")
+
+ >>> inputs = tokenizer(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt")
+ >>> with torch.inference_mode():
+ ... text_features = model.get_text_features(**inputs)
+ ```"""
+ text_outputs: BaseModelOutputWithPooling = self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ return_dict=True,
+ **kwargs,
+ )
+ pooled_output = text_outputs.pooler_output
+ text_outputs.pooler_output = self.text_projection(pooled_output)
+
+ return text_outputs
+
+ @can_return_tuple
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ Examples:
+
+ ```python
+ >>> import torch
+ >>> from transformers import AutoProcessor, GroupViTModel
+ >>> from transformers.image_utils import load_image
+
+ >>> model = GroupViTModel.from_pretrained("nvidia/groupvit-gcc-yfcc")
+ >>> processor = AutoProcessor.from_pretrained("nvidia/groupvit-gcc-yfcc")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> image = load_image(url)
+
+ >>> inputs = processor(images=image, return_tensors="pt")
+
+ >>> with torch.inference_mode():
+ ... image_features = model.get_image_features(**inputs)
+ ```"""
+ vision_outputs: BaseModelOutputWithPooling = self.vision_model(pixel_values, return_dict=True, **kwargs)
+ vision_outputs.pooler_output = self.visual_projection(vision_outputs.pooler_output)
+
+ return vision_outputs
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ return_loss: bool | None = None,
+ output_attentions: bool | None = None,
+ output_hidden_states: bool | None = None,
+ output_segmentation: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | GroupViTModelOutput:
+ r"""
+ return_loss (`bool`, *optional*):
+ Whether or not to return the contrastive loss.
+ output_segmentation (`bool`, *optional*):
+ Whether or not to return the segmentation logits.
+
+ Examples:
+
+ ```python
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+ >>> from transformers import AutoProcessor, GroupViTModel
+
+ >>> model = GroupViTModel.from_pretrained("nvidia/groupvit-gcc-yfcc")
+ >>> processor = AutoProcessor.from_pretrained("nvidia/groupvit-gcc-yfcc")
+
+ >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read()))
+
+ >>> inputs = processor(
+ ... text=["a photo of a cat", "a photo of a dog"], images=image, return_tensors="pt", padding=True
+ ... )
+
+ >>> outputs = model(**inputs)
+ >>> logits_per_image = outputs.logits_per_image # this is the image-text similarity score
+ >>> probs = logits_per_image.softmax(dim=1) # we can take the softmax to get the label probabilities
+ ```"""
+ # Use GROUPVIT model's config for some fields (if specified) instead of those of vision & text components.
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+ output_segmentation = (
+ output_segmentation if output_segmentation is not None else self.config.output_segmentation
+ )
+ if output_segmentation:
+ output_attentions = True
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+
+ # Vision side uses explicit flags (nn.Module-based, not hook-based)
+ vision_outputs = self.vision_model(
+ pixel_values=pixel_values,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=True,
+ )
+
+ text_outputs: BaseModelOutputWithPooling = self.text_model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ **kwargs,
+ )
+
+ image_embeds = vision_outputs.pooler_output
+ image_embeds = self.visual_projection(image_embeds)
+
+ text_embeds = text_outputs.pooler_output
+ text_embeds = self.text_projection(text_embeds)
+
+ # normalized features
+ image_embeds = image_embeds / image_embeds.norm(dim=-1, keepdim=True)
+ text_embeds = text_embeds / text_embeds.norm(dim=-1, keepdim=True)
+
+ # cosine similarity as logits
+ logit_scale = self.logit_scale.exp()
+ logits_per_text = torch.matmul(text_embeds, image_embeds.t()) * logit_scale
+ logits_per_image = logits_per_text.t()
+
+ seg_logits = None
+ if output_segmentation:
+ # grouped features
+ # [batch_size_image, num_group, hidden_size]
+ image_group_embeds = vision_outputs.last_hidden_state
+ # [batch_size_image*num_group, hidden_size]
+ image_group_embeds = self.visual_projection(image_group_embeds.reshape(-1, image_group_embeds.shape[-1]))
+ attentions = vision_outputs.attentions
+ # [batch_size_image, num_group, height, width]
+ grouping = get_grouping_from_attentions(attentions, pixel_values.shape[2:])
+
+ # normalized features
+ image_group_embeds = image_group_embeds / image_group_embeds.norm(dim=-1, keepdim=True)
+ # [batch_size_image x num_group, batch_size_text]
+ logits_per_image_group = torch.matmul(image_group_embeds, text_embeds.t()) * logit_scale
+ # [batch_size_image, batch_size_text, num_group]
+ logits_per_image_group = logits_per_image_group.reshape(
+ image_embeds.shape[0], -1, text_embeds.shape[0]
+ ).permute(0, 2, 1)
+
+ # [batch_size_image, batch_size_text, height x width]
+ flatten_grouping = grouping.reshape(grouping.shape[0], grouping.shape[1], -1)
+
+ # [batch_size_image, batch_size_text, height, width]
+ seg_logits = torch.matmul(logits_per_image_group, flatten_grouping) * logit_scale
+ seg_logits = seg_logits.reshape(
+ seg_logits.shape[0], seg_logits.shape[1], grouping.shape[2], grouping.shape[3]
+ )
+
+ loss = None
+ if return_loss:
+ loss = groupvit_loss(logits_per_text)
+
+ return GroupViTModelOutput(
+ loss=loss,
+ logits_per_image=logits_per_image,
+ logits_per_text=logits_per_text,
+ segmentation_logits=seg_logits,
+ text_embeds=text_embeds,
+ image_embeds=image_embeds,
+ text_model_output=text_outputs,
+ vision_model_output=vision_outputs,
+ )
+
+
+__all__ = ["GroupViTModel", "GroupViTPreTrainedModel", "GroupViTTextModel", "GroupViTVisionModel"]
diff --git a/third_party/transformers/src/transformers/models/hunyuan_v1_dense/__init__.py b/third_party/transformers/src/transformers/models/hunyuan_v1_dense/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..27de691c845369b97805fb53a8e509b7e948b386
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/hunyuan_v1_dense/__init__.py
@@ -0,0 +1,15 @@
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_hunyuan_v1_dense import *
+ from .modeling_hunyuan_v1_dense import *
+ from .tokenization_hy import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/hunyuan_v1_dense/configuration_hunyuan_v1_dense.py b/third_party/transformers/src/transformers/models/hunyuan_v1_dense/configuration_hunyuan_v1_dense.py
new file mode 100644
index 0000000000000000000000000000000000000000..7ac120cec14fb18ee2c86baafd4c360f370637d4
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/hunyuan_v1_dense/configuration_hunyuan_v1_dense.py
@@ -0,0 +1,64 @@
+# Copyright (C) 2025 THL A29 Limited, a Tencent company and the HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""HunYuanDenseV1 model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_rope_utils import RopeParameters
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="tencent/Hunyuan-7B-Instruct")
+@strict
+class HunYuanDenseV1Config(PreTrainedConfig):
+ r"""
+ eod_token_id (int, *optional*, defaults to 3):
+ Token ID representing the end-of-document marker. Used to indicate the termination of a text sequence.
+
+ Example: In multi-document processing, this token helps the model distinguish between separate documents.
+ """
+
+ model_type = "hunyuan_v1_dense"
+ keys_to_ignore_at_inference = ["past_key_values"]
+
+ vocab_size: int = 290943
+ hidden_size: int = 4096
+ intermediate_size: int = 11008
+ num_hidden_layers: int = 32
+ num_attention_heads: int = 32
+ num_key_value_heads: int | None = None
+ hidden_act: str = "silu"
+ max_position_embeddings: int = 2048
+ initializer_range: float = 0.02
+ rms_norm_eps: float = 1e-5
+ use_cache: bool = True
+ pad_token_id: int | None = 0
+ bos_token_id: int | None = 1
+ eos_token_id: int | list[int] | None = 2
+ eod_token_id: int | None = 3
+ pretraining_tp: int = 1
+ tie_word_embeddings: bool = False
+ rope_parameters: RopeParameters | dict | None = None
+ attention_bias: bool = False
+ attention_dropout: float | int = 0.0
+ head_dim: int | None = None
+
+ def __post_init__(self, **kwargs):
+ if self.num_key_value_heads is None:
+ self.num_key_value_heads = self.num_attention_heads
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["HunYuanDenseV1Config"]
diff --git a/third_party/transformers/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py b/third_party/transformers/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py
new file mode 100644
index 0000000000000000000000000000000000000000..d1652d78cbbc883993302d5eb768e297e6fbf616
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/hunyuan_v1_dense/modeling_hunyuan_v1_dense.py
@@ -0,0 +1,544 @@
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# This file was automatically generated from src/transformers/models/hunyuan_v1_dense/modular_hunyuan_v1_dense.py.
+# Do NOT edit this file manually as any edits will be overwritten by the generation of
+# the file from the modular. If any change should be done, please apply the change to the
+# modular_hunyuan_v1_dense.py file directly. One of our CI enforces this.
+# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright (C) 2025 THL A29 Limited, a Tencent company and the HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+
+from collections.abc import Callable
+from typing import Optional
+
+import torch
+from torch import nn
+
+from transformers.cache_utils import Cache
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import DynamicCache
+from ...generation import GenerationMixin
+from ...integrations import use_kernel_forward_from_hub, use_kernel_func_from_hub, use_kernelized_func
+from ...masking_utils import create_causal_mask
+from ...modeling_layers import GenericForSequenceClassification, GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple
+from ...utils.generic import maybe_autocast, merge_with_config_defaults
+from ...utils.output_capturing import capture_outputs
+from .configuration_hunyuan_v1_dense import HunYuanDenseV1Config
+
+
+@use_kernel_forward_from_hub("RMSNorm")
+class HunYuanDenseV1RMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps: float = 1e-6) -> None:
+ """
+ HunYuanDenseV1RMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ 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)
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+class HunYuanDenseV1MLP(nn.Module):
+ def __init__(self, config: HunYuanDenseV1Config, layer_idx=None, is_shared_mlp=False):
+ 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]
+ self.layer_idx = layer_idx
+
+ def forward(self, x):
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
+ return down_proj
+
+
+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)
+
+
+@use_kernel_func_from_hub("rotary_pos_emb")
+def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos.unsqueeze(unsqueeze_dim)
+ sin = sin.unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+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)
+
+
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs: Unpack[TransformersKwargs],
+):
+ key_states = repeat_kv(key, module.num_key_value_groups)
+ value_states = repeat_kv(value, module.num_key_value_groups)
+
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+ attn_output = torch.matmul(attn_weights, value_states)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+@use_kernelized_func(apply_rotary_pos_emb)
+class HunYuanDenseV1Attention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: HunYuanDenseV1Config, layer_idx: int):
+ super().__init__()
+ self.config = config
+ self.layer_idx = layer_idx
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
+ self.scaling = self.head_dim**-0.5
+ self.attention_dropout = config.attention_dropout
+ self.is_causal = True
+
+ self.q_proj = nn.Linear(
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.k_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.v_proj = nn.Linear(
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
+ )
+ self.o_proj = nn.Linear(
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
+ )
+ self.query_layernorm = HunYuanDenseV1RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+ self.key_layernorm = HunYuanDenseV1RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: torch.Tensor | None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+ query_states = self.query_layernorm(query_states)
+ key_states = self.key_layernorm(key_states)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class HunYuanDenseV1DecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: HunYuanDenseV1Config, layer_idx: int):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+
+ self.self_attn = HunYuanDenseV1Attention(config=config, layer_idx=layer_idx)
+
+ self.mlp = HunYuanDenseV1MLP(config)
+ self.input_layernorm = HunYuanDenseV1RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = HunYuanDenseV1RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.layer_idx = layer_idx
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ use_cache: bool | None = False,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = self.input_layernorm(hidden_states)
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ position_embeddings=position_embeddings,
+ **kwargs,
+ )
+ 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
+ return hidden_states
+
+
+@auto_docstring
+class HunYuanDenseV1PreTrainedModel(PreTrainedModel):
+ config: HunYuanDenseV1Config
+ base_model_prefix = "model"
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["HunYuanDenseV1DecoderLayer"]
+ _skip_keys_device_placement = ["past_key_values"]
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+
+ _can_compile_fullgraph = True
+ _supports_attention_backend = True
+ _can_record_outputs = {
+ "hidden_states": HunYuanDenseV1DecoderLayer,
+ "attentions": HunYuanDenseV1Attention,
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ super()._init_weights(module)
+
+ # DynamicNTKAlphaRotary - unique to this model
+ if "RotaryEmbedding" in module.__class__.__name__ and hasattr(module, "original_inv_freq"):
+ if module.rope_type == "dynamic" and module.config.rope_parameters.get("alpha"):
+ dim = module.config.head_dim
+ rope_theta = module.config.rope_parameters["rope_theta"]
+ alpha = module.config.rope_parameters["alpha"]
+
+ base = rope_theta * alpha ** (dim / (dim - 2))
+ buffer_value = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+ else:
+ rope_fn = (
+ ROPE_INIT_FUNCTIONS[module.rope_type]
+ if module.rope_type != "default"
+ else module.compute_default_rope_parameters
+ )
+ buffer_value, _ = rope_fn(module.config)
+ init.copy_(module.inv_freq, buffer_value)
+ init.copy_(module.original_inv_freq, buffer_value)
+
+
+class HunYuanDenseV1RotaryEmbedding(nn.Module):
+ inv_freq: torch.Tensor # fix linting for `register_buffer`
+
+ def __init__(self, config: HunYuanDenseV1Config, device=None):
+ super().__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+
+ # Diff from Llama - DynamicNTKAlphaRotary
+ if self.rope_type == "dynamic" and self.config.rope_parameters.get("alpha"):
+ self.dim = config.head_dim
+ base = self.config.rope_parameters["rope_theta"] * self.config.rope_parameters["alpha"] ** (
+ self.config.head_dim / (self.config.head_dim - 2)
+ )
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.config.head_dim))
+ self.attention_scaling = 1.0
+ else:
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+ @staticmethod
+ def compute_default_rope_parameters(
+ config: HunYuanDenseV1Config | None = None,
+ device: Optional["torch.device"] = None,
+ seq_len: int | None = None,
+ ) -> tuple["torch.Tensor", float]:
+ """
+ Computes the inverse frequencies according to the original RoPE implementation
+ Args:
+ config ([`~transformers.PreTrainedConfig`]):
+ The model configuration.
+ device (`torch.device`):
+ The device to use for initialization of the inverse frequencies.
+ seq_len (`int`, *optional*):
+ The current sequence length. Unused for this type of RoPE.
+ Returns:
+ Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the
+ post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE).
+ """
+ base = config.rope_parameters["rope_theta"]
+ dim = getattr(config, "head_dim", None) or config.hidden_size // config.num_attention_heads
+
+ attention_factor = 1.0 # Unused in this type of RoPE
+
+ # Compute the inverse frequencies
+ inv_freq = 1.0 / (
+ base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(device=device, dtype=torch.float) / dim)
+ )
+ return inv_freq, attention_factor
+
+ @torch.no_grad()
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
+ def forward(self, x, position_ids):
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
+ position_ids_expanded = position_ids[:, None, :].float()
+
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
+ with maybe_autocast(device_type=device_type, enabled=False): # Force float32
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
+
+
+@auto_docstring
+class HunYuanDenseV1Model(HunYuanDenseV1PreTrainedModel):
+ def __init__(self, config: HunYuanDenseV1Config):
+ super().__init__(config)
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+ self.layers = nn.ModuleList(
+ [HunYuanDenseV1DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.norm = HunYuanDenseV1RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.rotary_emb = HunYuanDenseV1RotaryEmbedding(config=config)
+ self.gradient_checkpointing = False
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ use_cache: bool | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> BaseModelOutputWithPast:
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds: torch.Tensor = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ if position_ids is None:
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
+ position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + past_seen_tokens
+ position_ids = position_ids.unsqueeze(0)
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ hidden_states = inputs_embeds
+ position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+
+ for decoder_layer in self.layers[: self.config.num_hidden_layers]:
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_embeddings=position_embeddings,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+ return BaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+@auto_docstring
+class HunYuanDenseV1ForCausalLM(HunYuanDenseV1PreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+ _tp_plan = {"lm_head": "colwise_gather_output"}
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
+
+ def __init__(self, config):
+ super().__init__(config)
+ self.model = HunYuanDenseV1Model(config)
+ self.vocab_size = config.vocab_size
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> CausalLMOutputWithPast:
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import AutoTokenizer, HunYuanDenseV1ForCausalLM
+
+ >>> model = HunYuanDenseV1ForCausalLM.from_pretrained("meta-hunyuan_v1_dense/HunYuanDenseV1-2-7b-hf")
+ >>> tokenizer = AutoTokenizer.from_pretrained("meta-hunyuan_v1_dense/HunYuanDenseV1-2-7b-hf")
+
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
+
+ >>> # Generate
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
+ ```"""
+ outputs: BaseModelOutputWithPast = self.model(
+ input_ids=input_ids,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ inputs_embeds=inputs_embeds,
+ use_cache=use_cache,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return CausalLMOutputWithPast(
+ loss=loss,
+ logits=logits,
+ past_key_values=outputs.past_key_values,
+ hidden_states=outputs.hidden_states,
+ attentions=outputs.attentions,
+ )
+
+
+class HunYuanDenseV1ForSequenceClassification(GenericForSequenceClassification, HunYuanDenseV1PreTrainedModel):
+ pass
+
+
+__all__ = [
+ "HunYuanDenseV1ForCausalLM",
+ "HunYuanDenseV1Model",
+ "HunYuanDenseV1PreTrainedModel",
+ "HunYuanDenseV1ForSequenceClassification",
+]
diff --git a/third_party/transformers/src/transformers/models/hunyuan_v1_dense/modular_hunyuan_v1_dense.py b/third_party/transformers/src/transformers/models/hunyuan_v1_dense/modular_hunyuan_v1_dense.py
new file mode 100644
index 0000000000000000000000000000000000000000..cccad38a7a14b7f98752a9cd6f03ad3c13babe66
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/hunyuan_v1_dense/modular_hunyuan_v1_dense.py
@@ -0,0 +1,188 @@
+# Copyright (C) 2025 THL A29 Limited, a Tencent company and the HuggingFace Inc. team. All rights reserved.
+#
+# 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 HunYuanDenseV1 model."""
+
+from collections.abc import Callable
+
+import torch
+from torch import nn
+
+from transformers.cache_utils import Cache
+from transformers.utils import (
+ logging,
+)
+
+from ... import initialization as init
+from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs
+from ..llama.modeling_llama import (
+ LlamaAttention,
+ LlamaDecoderLayer,
+ LlamaForCausalLM,
+ LlamaForSequenceClassification,
+ LlamaMLP,
+ LlamaModel,
+ LlamaPreTrainedModel,
+ LlamaRMSNorm,
+ LlamaRotaryEmbedding,
+ apply_rotary_pos_emb,
+ eager_attention_forward,
+)
+from .configuration_hunyuan_v1_dense import HunYuanDenseV1Config
+
+
+logger = logging.get_logger(__name__)
+
+
+class HunYuanDenseV1RMSNorm(LlamaRMSNorm):
+ pass
+
+
+class HunYuanDenseV1MLP(LlamaMLP):
+ def __init__(self, config: HunYuanDenseV1Config, layer_idx=None, is_shared_mlp=False):
+ super().__init__(config)
+ self.layer_idx = layer_idx
+ 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)
+
+
+class HunYuanDenseV1Attention(LlamaAttention):
+ def __init__(self, config: HunYuanDenseV1Config, layer_idx: int):
+ super().__init__(config, layer_idx)
+ self.query_layernorm = HunYuanDenseV1RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+ self.key_layernorm = HunYuanDenseV1RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
+ attention_mask: torch.Tensor | None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ input_shape = hidden_states.shape[:-1]
+ hidden_shape = (*input_shape, -1, self.head_dim)
+
+ query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+ cos, sin = position_embeddings
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+ query_states = self.query_layernorm(query_states)
+ key_states = self.key_layernorm(key_states)
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+ return attn_output, attn_weights
+
+
+class HunYuanDenseV1DecoderLayer(LlamaDecoderLayer):
+ def __init__(self, config: HunYuanDenseV1Config, layer_idx: int):
+ super().__init__(config, layer_idx)
+ self.layer_idx = layer_idx
+
+
+class HunYuanDenseV1PreTrainedModel(LlamaPreTrainedModel, PreTrainedModel):
+ @torch.no_grad()
+ def _init_weights(self, module):
+ PreTrainedModel._init_weights(self, module)
+
+ # DynamicNTKAlphaRotary - unique to this model
+ if "RotaryEmbedding" in module.__class__.__name__ and hasattr(module, "original_inv_freq"):
+ if module.rope_type == "dynamic" and module.config.rope_parameters.get("alpha"):
+ dim = module.config.head_dim
+ rope_theta = module.config.rope_parameters["rope_theta"]
+ alpha = module.config.rope_parameters["alpha"]
+
+ base = rope_theta * alpha ** (dim / (dim - 2))
+ buffer_value = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+ else:
+ rope_fn = (
+ ROPE_INIT_FUNCTIONS[module.rope_type]
+ if module.rope_type != "default"
+ else module.compute_default_rope_parameters
+ )
+ buffer_value, _ = rope_fn(module.config)
+ init.copy_(module.inv_freq, buffer_value)
+ init.copy_(module.original_inv_freq, buffer_value)
+
+
+class HunYuanDenseV1RotaryEmbedding(LlamaRotaryEmbedding):
+ def __init__(self, config: HunYuanDenseV1Config, device=None):
+ nn.Module.__init__()
+ self.max_seq_len_cached = config.max_position_embeddings
+ self.original_max_seq_len = config.max_position_embeddings
+
+ self.config = config
+
+ self.rope_type = self.config.rope_parameters["rope_type"]
+
+ # Diff from Llama - DynamicNTKAlphaRotary
+ if self.rope_type == "dynamic" and self.config.rope_parameters.get("alpha"):
+ self.dim = config.head_dim
+ base = self.config.rope_parameters["rope_theta"] * self.config.rope_parameters["alpha"] ** (
+ self.config.head_dim / (self.config.head_dim - 2)
+ )
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.config.head_dim))
+ self.attention_scaling = 1.0
+ else:
+ rope_init_fn: Callable = self.compute_default_rope_parameters
+ if self.rope_type != "default":
+ rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
+ inv_freq, self.attention_scaling = rope_init_fn(self.config, device)
+
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.register_buffer("original_inv_freq", inv_freq.clone(), persistent=False)
+
+
+class HunYuanDenseV1Model(LlamaModel):
+ pass
+
+
+class HunYuanDenseV1ForCausalLM(LlamaForCausalLM):
+ pass
+
+
+class HunYuanDenseV1ForSequenceClassification(LlamaForSequenceClassification):
+ pass
+
+
+__all__ = [
+ "HunYuanDenseV1ForCausalLM",
+ "HunYuanDenseV1Model",
+ "HunYuanDenseV1PreTrainedModel",
+ "HunYuanDenseV1ForSequenceClassification",
+]
diff --git a/third_party/transformers/src/transformers/models/idefics/__init__.py b/third_party/transformers/src/transformers/models/idefics/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..5b92e53c1f81735e24b7273844c4a20775cecaa8
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/idefics/__init__.py
@@ -0,0 +1,30 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_idefics import *
+ from .image_processing_idefics import *
+ from .image_processing_pil_idefics import *
+ from .modeling_idefics import *
+ from .processing_idefics import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/idefics/configuration_idefics.py b/third_party/transformers/src/transformers/models/idefics/configuration_idefics.py
new file mode 100644
index 0000000000000000000000000000000000000000..21e3860022ef5f6bd26d376fddc12c1d294b9e96
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/idefics/configuration_idefics.py
@@ -0,0 +1,169 @@
+# Copyright 2022 EleutherAI 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.
+"""Idefics model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+
+
+@auto_docstring(checkpoint="HuggingFaceM4/idefics-9b")
+@strict
+class IdeficsVisionConfig(PreTrainedConfig):
+ model_type = "idefics_vision"
+ attribute_map = {"hidden_size": "embed_dim"}
+
+ embed_dim: int = 768
+ image_size: int | list[int] | tuple[int, int] = 224
+ intermediate_size: int = 5120
+ patch_size: int | list[int] | tuple[int, int] = 14
+ num_hidden_layers: int = 32
+ num_attention_heads: int = 16
+ num_channels: int = 3
+ hidden_act: str = "gelu"
+ layer_norm_eps: float = 1e-5
+ attention_dropout: float | int = 0.0
+ initializer_range: float = 0.02
+ initializer_factor: float = 1.0
+
+
+@auto_docstring(checkpoint="HuggingFaceM4/idefics-9b")
+@strict
+class IdeficsPerceiverConfig(PreTrainedConfig):
+ r"""
+ use_resampler (`bool`, *optional*, defaults to `False`):
+ Whether or not to use the resampler
+ resampler_n_latents (`int`, *optional*, defaults to 64):
+ Number of latent embeddings to resample ("compress") the input sequence to (usually < 128).
+ resampler_depth (`int`, *optional*, defaults to 6):
+ Depth of the Perceiver Resampler (Transformer w/ cross attention). Should be shallow (< 3).
+ resampler_n_heads (`int`, *optional*, defaults to 16):
+ Number of heads in each Transformer block (for multi-headed self-attention).
+ resampler_head_dim (`int`, *optional*, defaults to 96):
+ Dimensionality of each head projection in the Transformer block.
+ qk_layer_norms_perceiver (`bool`, *optional*, defaults to `False`):
+ Whether or not to use qk layer norms in perceiver
+ """
+
+ model_type = "idefics_perciever"
+
+ use_resampler: bool = False
+ resampler_n_latents: int = 64
+ resampler_depth: int = 6
+ resampler_n_heads: int = 16
+ resampler_head_dim: int = 96
+ qk_layer_norms_perceiver: bool = False
+
+
+@auto_docstring(checkpoint="HuggingFaceM4/idefics-9b")
+@strict
+class IdeficsConfig(PreTrainedConfig):
+ r"""
+ additional_vocab_size (`int`, *optional*, defaults to 0):
+ Additional vocabulary size of the model, typically for the special "![]()
" token. Additional vocab tokens
+ are always trainable whereas regular vocab tokens can be frozen or not.
+ alpha_initializer (`str`, *optional*, defaults to `"zeros"`):
+ Initialization type for the alphas.
+ alphas_initializer_range (`float`, *optional*, defaults to 0.0):
+ The standard deviation of the truncated_normal_initializer for initializing the alphas in the Gated Cross
+ Attention.
+ alpha_type (`str`, *optional*, defaults to `"float"`):
+ Whether the gating alphas should be vectors or single floats.
+ cross_layer_interval (`int`, *optional*, default to 1):
+ Interval for cross attention (from text to image) layers.
+ qk_layer_norms (`bool`, *optional*, defaults to `False`):
+ Whether to add layer norm after q and k
+ freeze_text_layers (`bool`, *optional*, defaults to `True`):
+ Whether to freeze text layers
+ freeze_text_module_exceptions (`bool`, *optional*, defaults to `[]`):
+ Exceptions to freezing text layers when `freeze_text_layers` is `True`
+ freeze_lm_head (`bool`, *optional*, defaults to `False`):
+ Whether to freeze lm head
+ freeze_vision_layers (`bool`, *optional*, defaults to `True`):
+ Whether to freeze vision layers
+ freeze_vision_module_exceptions (`bool`, *optional*, defaults to `[]`):
+ Exceptions to freezing vision layers when `freeze_vision_layers` is `True`
+ use_resampler (`bool`, *optional*, defaults to `False`):
+ Whether to use the Resampler
+ perceiver_config (`IdeficsPerceiverConfig`, *optional*):
+ Custom perceiver config or dict
+
+ Example:
+
+ ```python
+ >>> from transformers import IdeficsModel, IdeficsConfig
+
+ >>> # Initializing a Idefics idefics-9b style configuration
+ >>> configuration = IdeficsConfig()
+
+ >>> # Initializing a model from the idefics-9b style configuration
+ >>> model = IdeficsModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "idefics"
+ sub_configs = {"perceiver_config": IdeficsPerceiverConfig, "vision_config": IdeficsVisionConfig}
+
+ vocab_size: int = 32000
+ additional_vocab_size: int = 0
+ hidden_size: int = 4096
+ intermediate_size: int = 11008
+ num_hidden_layers: int = 32
+ num_attention_heads: int = 32
+ dropout: float | int = 0.0
+ hidden_act: str = "silu"
+ initializer_range: float = 0.02
+ alpha_initializer: str = "zeros"
+ alphas_initializer_range: float = 0.0
+ alpha_type: str = "float"
+ rms_norm_eps: float = 1e-6
+ use_cache: bool = True
+ pad_token_id: int | None = 0
+ bos_token_id: int | None = 1
+ eos_token_id: int | list[int] | None = 2
+ tie_word_embeddings: bool = False
+ cross_layer_interval: int = 1
+ qk_layer_norms: bool = False
+ freeze_text_layers: bool = True
+ freeze_text_module_exceptions: list | tuple = ()
+ freeze_lm_head: bool = False
+ freeze_vision_layers: bool = True
+ freeze_vision_module_exceptions: list | tuple = ()
+ use_resampler: bool = False
+ vision_config: dict | PreTrainedConfig | None = None
+ perceiver_config: dict | PreTrainedConfig | None = None
+
+ def __post_init__(self, **kwargs):
+ if self.perceiver_config is None:
+ self.perceiver_config = IdeficsPerceiverConfig()
+ elif isinstance(self.perceiver_config, dict):
+ self.perceiver_config = IdeficsPerceiverConfig(**self.perceiver_config)
+
+ if self.vision_config is None:
+ self.vision_config = IdeficsVisionConfig()
+ elif isinstance(self.vision_config, dict):
+ self.vision_config = IdeficsVisionConfig(**self.vision_config)
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["IdeficsConfig"]
diff --git a/third_party/transformers/src/transformers/models/idefics/image_processing_idefics.py b/third_party/transformers/src/transformers/models/idefics/image_processing_idefics.py
new file mode 100644
index 0000000000000000000000000000000000000000..ce6f270329712dcd756f24fa8d5684b464ef1346
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/idefics/image_processing_idefics.py
@@ -0,0 +1,90 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Image processor class for Idefics."""
+
+from collections.abc import Callable
+
+from ...image_processing_backends import TorchvisionBackend
+from ...image_processing_utils import BatchFeature
+from ...image_utils import (
+ ImageInput,
+ PILImageResampling,
+ make_flat_list_of_images,
+)
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import TensorType, auto_docstring, is_torch_available
+
+
+IDEFICS_STANDARD_MEAN = [0.48145466, 0.4578275, 0.40821073]
+IDEFICS_STANDARD_STD = [0.26862954, 0.26130258, 0.27577711]
+
+
+class IdeficsImageProcessorKwargs(ImagesKwargs, total=False):
+ r"""
+ transform (`Callable`, *optional*, defaults to `None`):
+ A custom transform function that accepts a single image can be passed for training. For example,
+ `torchvision.Compose` can be used to compose multiple transforms. If `None` - an inference mode is
+ assumed - and then a preset of inference-specific transforms will be applied to the images.
+ image_size (`int`, *optional*, defaults to `self.image_size`):
+ Resize to image size. This is a backward-compatible alias for `size`. When provided, it overrides
+ `size` and sets it to `{"height": image_size, "width": image_size}`.
+ image_num_channels (`int`, *optional*, defaults to `3`):
+ The number of channels of the image.
+ """
+
+ transform: Callable | None
+ image_size: int
+ image_num_channels: int
+
+
+@auto_docstring
+class IdeficsImageProcessor(TorchvisionBackend):
+ valid_kwargs = IdeficsImageProcessorKwargs
+ resample = PILImageResampling.BICUBIC
+ image_mean = IDEFICS_STANDARD_MEAN
+ image_std = IDEFICS_STANDARD_STD
+ size = {"height": 224, "width": 224}
+ do_resize = True
+ do_rescale = True
+ do_normalize = True
+ do_convert_rgb = True
+ image_num_channels = 3
+
+ def __init__(self, **kwargs: Unpack[IdeficsImageProcessorKwargs]):
+ image_size = kwargs.pop("image_size", None)
+ if image_size is not None:
+ kwargs["size"] = {"height": image_size, "width": image_size}
+ super().__init__(**kwargs)
+ self.image_size = self.size.height
+
+ @auto_docstring
+ def preprocess(
+ self,
+ images: ImageInput,
+ **kwargs: Unpack[IdeficsImageProcessorKwargs],
+ ) -> "TensorType | BatchFeature":
+ transform = kwargs.pop("transform", None)
+ if transform is not None:
+ if not is_torch_available():
+ raise ImportError("To pass in `transform` torch must be installed")
+ import torch
+
+ images = self.fetch_images(images)
+ images = make_flat_list_of_images(images)
+ images = [transform(x) for x in images]
+ return torch.stack(images)
+ return super().preprocess(images, **kwargs).pixel_values
+
+
+__all__ = ["IdeficsImageProcessor"]
diff --git a/third_party/transformers/src/transformers/models/idefics/image_processing_pil_idefics.py b/third_party/transformers/src/transformers/models/idefics/image_processing_pil_idefics.py
new file mode 100644
index 0000000000000000000000000000000000000000..d74ade839489833eb0ee52a968b9ab44be354112
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/idefics/image_processing_pil_idefics.py
@@ -0,0 +1,94 @@
+# Copyright 2022 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""Image processor class for Idefics."""
+
+from collections.abc import Callable
+
+from ...image_processing_backends import PilBackend
+from ...image_utils import ImageInput, PILImageResampling, make_flat_list_of_images
+from ...processing_utils import ImagesKwargs, Unpack
+from ...utils import auto_docstring, is_torch_available
+
+
+# Adapted from transformers.models.idefics.image_processing_idefics.IDEFICS_STANDARD_MEAN
+IDEFICS_STANDARD_MEAN = [0.48145466, 0.4578275, 0.40821073]
+
+# Adapted from transformers.models.idefics.image_processing_idefics.IDEFICS_STANDARD_STD
+IDEFICS_STANDARD_STD = [0.26862954, 0.26130258, 0.27577711]
+
+
+# Adapted from transformers.models.idefics.image_processing_idefics.IdeficsImageProcessorKwargs
+class IdeficsImageProcessorKwargs(ImagesKwargs, total=False):
+ r"""
+ transform (`Callable`, *optional*, defaults to `None`):
+ A custom transform function that accepts a single image can be passed for training. For example,
+ `torchvision.Compose` can be used to compose multiple transforms. If `None` - an inference mode is
+ assumed - and then a preset of inference-specific transforms will be applied to the images.
+ image_size (`int`, *optional*, defaults to `self.image_size`):
+ Resize to image size. This is a backward-compatible alias for `size`. When provided, it overrides
+ `size` and sets it to `{"height": image_size, "width": image_size}`.
+ image_num_channels (`int`, *optional*, defaults to `3`):
+ The number of channels of the image.
+ """
+
+ transform: Callable | None
+ image_size: int
+ image_num_channels: int
+
+
+@auto_docstring
+class IdeficsImageProcessorPil(PilBackend):
+ valid_kwargs = IdeficsImageProcessorKwargs
+ resample = PILImageResampling.BICUBIC
+ image_mean = IDEFICS_STANDARD_MEAN
+ image_std = IDEFICS_STANDARD_STD
+ size = {"height": 224, "width": 224}
+ do_resize = True
+ do_rescale = True
+ do_normalize = True
+ do_convert_rgb = True
+ image_num_channels = 3
+
+ def __init__(self, **kwargs: Unpack[IdeficsImageProcessorKwargs]):
+ image_size = kwargs.pop("image_size", None)
+ if image_size is not None:
+ kwargs["size"] = {"height": image_size, "width": image_size}
+ super().__init__(**kwargs)
+ self.image_size = self.size.height
+
+ @auto_docstring
+ def preprocess(
+ self,
+ images: ImageInput,
+ **kwargs: Unpack[IdeficsImageProcessorKwargs],
+ ):
+ r"""
+ transform (`Callable`, *optional*, defaults to `None`):
+ A custom transform function that accepts a single image can be passed for training. If `None`,
+ inference-mode transforms are applied.
+ """
+ transform = kwargs.pop("transform", None)
+ if transform is not None:
+ if not is_torch_available():
+ raise ImportError("To pass in `transform` torch must be installed")
+ import torch
+
+ images = self.fetch_images(images)
+ images = make_flat_list_of_images(images)
+ images = [transform(x) for x in images]
+ return torch.stack(images)
+ return super().preprocess(images, **kwargs).pixel_values
+
+
+__all__ = ["IdeficsImageProcessorPil"]
diff --git a/third_party/transformers/src/transformers/models/idefics/modeling_idefics.py b/third_party/transformers/src/transformers/models/idefics/modeling_idefics.py
new file mode 100644
index 0000000000000000000000000000000000000000..8f4578e1d0f27082d9f80e23da3f9d7a73d9e706
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/idefics/modeling_idefics.py
@@ -0,0 +1,1277 @@
+# Copyright 2022 EleutherAI 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 Idefics model."""
+
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...cache_utils import Cache, DynamicCache
+from ...generation import GenerationMixin
+from ...masking_utils import create_causal_mask
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import ModelOutput
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedConfig, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from .configuration_idefics import IdeficsConfig
+from .perceiver import IdeficsPerceiverResampler
+from .vision import IdeficsVisionEmbeddings, IdeficsVisionTransformer
+
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for Idefics model's outputs that may also contain a past key/values (to speed up sequential decoding).
+ """
+)
+class IdeficsBaseModelOutputWithPast(ModelOutput):
+ r"""
+ 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 (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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.
+ 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 = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ image_hidden_states: tuple[torch.FloatTensor] | None = None
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Base class for Idefics causal language model (or autoregressive) outputs.
+ """
+)
+class IdeficsCausalLMOutputWithPast(ModelOutput):
+ r"""
+ 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 (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):
+ It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).
+
+ 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.
+ 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: torch.FloatTensor | None = None
+ logits: torch.FloatTensor | None = None
+ past_key_values: Cache | None = None
+ hidden_states: tuple[torch.FloatTensor] | None = None
+ attentions: tuple[torch.FloatTensor] | None = None
+ image_hidden_states: tuple[torch.FloatTensor] | None = 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")
+ model_kwargs["image_encoder_embeddings"] = model_kwargs.get("image_encoder_embeddings")
+ model_kwargs["perceiver_embeddings"] = model_kwargs.get("perceiver_embeddings")
+ model_kwargs["image_attention_mask"] = model_kwargs.get("image_attention_mask")
+
+ 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["image_attention_mask"] is not None:
+ model_kwargs["image_attention_mask"] = model_kwargs["image_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_encoder_embeddings"] is not None:
+ model_kwargs["image_encoder_embeddings"] = model_kwargs["image_encoder_embeddings"].index_select(
+ 0, expanded_return_idx
+ )
+
+ elif model_kwargs["perceiver_embeddings"] is not None:
+ model_kwargs["perceiver_embeddings"] = model_kwargs["perceiver_embeddings"].index_select(
+ 0, expanded_return_idx
+ )
+
+ return input_ids, model_kwargs
+
+
+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 IdeficsDecoupledEmbedding(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: bool | None = False,
+ device=None,
+ dtype=None,
+ padding_idx=None,
+ **kwargs,
+ ) -> None:
+ """
+ Args:
+ num_embeddings (`int`):
+ Size of the dictionary of embeddings
+ num_additional_embeddings (`int`):
+ Number of additional embeddings. Only useful when you `partially_freeze=True`.
+ embedding_dim (`int`):
+ The size of each embedding vector
+ partially_freeze: (`bool`, *optional*, defaults to `False`):
+ If `True`, the regular `weight` will be frozen. `additional_weight` is never frozen.
+ padding_idx (`int`, *optional*):
+ The padding index (needs to be less than num_embeddings)
+
+ 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 F.embedding(input_ids, self.weight)
+
+ # 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 f"num_embeddings={self.num_embeddings}, num_additional_embeddings={self.num_additional_embeddings}, embedding_dim={self.embedding_dim}, partially_freeze={self.partially_freeze}"
+
+
+class IdeficsDecoupledLinear(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 f"in_features={self.in_features}, out_features={self.out_features}, out_additional_features={self.out_additional_features}, bias={self.bias is not None}, partially_freeze={self.partially_freeze}"
+
+
+# this was adapted from LlamaRMSNorm
+class IdeficsRMSNorm(nn.Module):
+ def __init__(self, hidden_size, eps=1e-6):
+ """
+ IdeficsRMSNorm is equivalent to T5LayerNorm
+ """
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(hidden_size))
+ self.variance_epsilon = eps
+
+ def forward(self, hidden_states):
+ variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+
+ # convert into half-precision if necessary
+ if self.weight.dtype in [torch.float16, torch.bfloat16]:
+ hidden_states = hidden_states.to(self.weight.dtype)
+
+ return self.weight * hidden_states
+
+ def extra_repr(self):
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
+
+
+# this was adapted from LlamaRotaryEmbedding
+class IdeficsEmbedding(torch.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, dtype=torch.int64).to(device=device, dtype=torch.float) / 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=torch.int64).type_as(self.inv_freq)
+
+ 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),
+ )
+
+
+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)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
+ """Applies Rotary Position Embedding to the query and key tensors.
+
+ Args:
+ q (`torch.Tensor`): The query tensor.
+ k (`torch.Tensor`): The key tensor.
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
+ position_ids (`torch.Tensor`):
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
+ used to pass offsetted position ids when working with a KV-cache.
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
+ Returns:
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
+ """
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim)
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim)
+ q_embed = (q * cos) + (rotate_half(q) * sin)
+ k_embed = (k * cos) + (rotate_half(k) * sin)
+ return q_embed, k_embed
+
+
+# this was adapted from LlamaMLP
+class IdeficsMLP(nn.Module):
+ def __init__(
+ self,
+ hidden_size: int,
+ intermediate_size: int,
+ hidden_act: str,
+ ):
+ super().__init__()
+ self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+ self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+ self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+ self.act_fn = ACT2FN[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.siglip.modeling_siglip.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+# this was adapted from LlamaAttention
+class IdeficsAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(
+ self,
+ hidden_size: int,
+ num_heads: int,
+ dropout: float = 0.0,
+ is_cross_attention: bool = False,
+ config: PreTrainedConfig | None = None,
+ qk_layer_norms: bool = False,
+ layer_idx: int | None = None,
+ ):
+ super().__init__()
+ self.config = config
+ self.hidden_size = hidden_size
+ self.num_heads = num_heads
+ self.head_dim = hidden_size // num_heads
+ self.dropout = dropout
+ self.is_causal = True
+ self.scaling = self.head_dim**-0.5
+
+ self.layer_idx = layer_idx
+ if layer_idx is None:
+ logger.warning_once(
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
+ "when creating this class."
+ )
+
+ if (self.head_dim * 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`: {num_heads})."
+ )
+
+ self.is_cross_attention = is_cross_attention
+
+ if not hasattr(nn.functional, "scaled_dot_product_attention"):
+ raise ValueError("this model requires pytorch 2.0 or higher")
+
+ if self.is_cross_attention:
+ kv_input_dim = (
+ self.hidden_size if not hasattr(config.vision_config, "embed_dim") else config.vision_config.embed_dim
+ )
+ self.q_proj = nn.Linear(
+ self.hidden_size,
+ num_heads * self.head_dim,
+ bias=False,
+ )
+ self.k_proj = nn.Linear(kv_input_dim, num_heads * self.head_dim, bias=False)
+ self.v_proj = nn.Linear(
+ kv_input_dim,
+ num_heads * self.head_dim,
+ bias=False,
+ )
+ else:
+ self.q_proj = nn.Linear(
+ self.hidden_size,
+ num_heads * self.head_dim,
+ bias=False,
+ )
+ self.k_proj = nn.Linear(
+ self.hidden_size,
+ num_heads * self.head_dim,
+ bias=False,
+ )
+ self.v_proj = nn.Linear(
+ self.hidden_size,
+ num_heads * self.head_dim,
+ bias=False,
+ )
+ self.o_proj = nn.Linear(
+ num_heads * self.head_dim,
+ hidden_size,
+ bias=False,
+ )
+ self.rotary_emb = IdeficsEmbedding(self.head_dim)
+
+ self.qk_layer_norms = qk_layer_norms
+ if self.qk_layer_norms:
+ self.q_layer_norm = IdeficsRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+ self.k_layer_norm = IdeficsRMSNorm(self.head_dim, eps=config.rms_norm_eps)
+
+ 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: torch.Tensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ # if key_value_states are provided this layer is used as a cross-attention layer
+ is_cross_attention = self.is_cross_attention or key_value_states is not None
+
+ 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)
+ if not is_cross_attention:
+ key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+ value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+ else:
+ _, kv_len, _ = key_value_states.size() # Note that, in this case, `kv_len` == `kv_seq_len`
+ key_states = self.k_proj(key_value_states).view(bsz, kv_len, self.num_heads, self.head_dim).transpose(1, 2)
+ value_states = (
+ self.v_proj(key_value_states).view(bsz, kv_len, self.num_heads, self.head_dim).transpose(1, 2)
+ )
+
+ kv_seq_len = key_states.shape[-2]
+ if past_key_values is not None:
+ kv_seq_len += past_key_values.get_seq_length()
+
+ if not is_cross_attention:
+ cos, sin = self.rotary_emb(value_states, seq_len=max(kv_seq_len, q_len))
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+ # [bsz, nh, t, hd]
+
+ if past_key_values is not None:
+ key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
+
+ if self.qk_layer_norms:
+ query_states = self.q_layer_norm(query_states)
+ key_states = self.k_layer_norm(key_states)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask,
+ dropout=0.0 if not self.training else self.dropout,
+ scaling=self.scaling,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
+ attn_output = self.o_proj(attn_output)
+
+ return attn_output, attn_weights
+
+
+# this was adapted from LlamaDecoderLayer
+class IdeficsDecoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: IdeficsConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.self_attn = IdeficsAttention(
+ hidden_size=self.hidden_size,
+ num_heads=config.num_attention_heads,
+ dropout=config.dropout,
+ config=config,
+ layer_idx=layer_idx,
+ )
+ self.mlp = IdeficsMLP(
+ hidden_size=self.hidden_size,
+ intermediate_size=config.intermediate_size,
+ hidden_act=config.hidden_act,
+ )
+ self.input_layernorm = IdeficsRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = IdeficsRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.dropout = config.dropout
+
+ @auto_docstring
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.FloatTensor:
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ 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 = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+class IdeficsGatedCrossAttentionLayer(GradientCheckpointingLayer):
+ def __init__(self, config: IdeficsConfig, layer_idx: int | None = None):
+ super().__init__()
+ self.hidden_size = config.hidden_size
+ self.cross_attn = IdeficsAttention(
+ hidden_size=self.hidden_size,
+ num_heads=config.num_attention_heads,
+ is_cross_attention=True,
+ dropout=config.dropout,
+ config=config,
+ qk_layer_norms=config.qk_layer_norms,
+ layer_idx=layer_idx,
+ )
+ self.mlp = IdeficsMLP(
+ hidden_size=self.hidden_size,
+ intermediate_size=config.intermediate_size,
+ hidden_act=config.hidden_act,
+ )
+ self.input_layernorm = IdeficsRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = IdeficsRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.config = config.dropout
+
+ self.act_cross_attn = nn.Tanh()
+ self.act_dense = nn.Tanh()
+
+ if config.alpha_initializer == "zeros":
+ if config.alpha_type == "vector":
+ self.alpha_cross_attn = nn.Parameter(torch.zeros(1, 1, self.hidden_size))
+ self.alpha_dense = nn.Parameter(torch.zeros(1, 1, self.hidden_size))
+ elif config.alpha_type == "float":
+ self.alpha_cross_attn = nn.Parameter(torch.zeros(1))
+ self.alpha_dense = nn.Parameter(torch.zeros(1))
+ else:
+ raise ValueError(f"Unknown value for `alpha_type` ({config.alpha_type})")
+
+ elif config.alpha_initializer == "ones":
+ if config.alpha_type == "vector":
+ self.alpha_cross_attn = nn.Parameter(torch.ones(1, 1, self.hidden_size))
+ self.alpha_dense = nn.Parameter(torch.ones(1, 1, self.hidden_size))
+ elif config.alpha_type == "float":
+ self.alpha_cross_attn = nn.Parameter(torch.ones(1))
+ self.alpha_dense = nn.Parameter(torch.ones(1))
+ else:
+ raise ValueError(f"Unknown value for `alpha_type` ({config.alpha_type})")
+
+ elif config.alpha_initializer in {"normal", "gaussian", "random"}:
+ if config.alpha_type == "vector":
+ self.alpha_cross_attn = nn.Parameter(
+ torch.normal(mean=0.0, std=config.alphas_initializer_range, size=(1, 1, self.hidden_size))
+ )
+ self.alpha_dense = nn.Parameter(
+ torch.normal(mean=0.0, std=config.alphas_initializer_range, size=(1, 1, self.hidden_size))
+ )
+ elif config.alpha_type == "float":
+ self.alpha_cross_attn = nn.Parameter(
+ torch.normal(mean=0.0, std=config.alphas_initializer_range, size=(1))
+ )
+ self.alpha_dense = nn.Parameter(torch.normal(mean=0.0, std=config.alphas_initializer_range, size=(1)))
+ else:
+ raise ValueError(f"Unknown value for `alpha_type` ({config.alpha_type})")
+
+ else:
+ raise NotImplementedError(f"Alpha initialization scheme {config.alpha_initializer} not yet implemented!")
+
+ if not (hasattr(self, "alpha_cross_attn") and hasattr(self, "alpha_dense")):
+ raise ValueError("Alpha parameters not initialized correctly!")
+
+ @auto_docstring
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ image_hidden_states: torch.Tensor | None = None,
+ image_attention_mask: torch.Tensor | None = None,
+ cross_attention_gate: torch.Tensor | None = None,
+ past_key_values: Cache | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.FloatTensor:
+ r"""
+ image_hidden_states (`torch.FloatTensor`):
+ Input to the layer of shape `(batch, seq_len, embed_dim)`
+ image_attention_mask (`torch.FloatTensor`, *optional*):
+ image attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ cross_attention_gate (`torch.FloatTensor`, *optional*):
+ gate of size `(batch, seq_len)` used to zero-out cross-attention output for tokens attending no images.
+ """
+ if image_hidden_states is None:
+ raise ValueError(
+ "`image_hidden_states` is required for Idefics cross attention module which are visual features to be"
+ " conditioned on."
+ )
+
+ if cross_attention_gate is None:
+ raise ValueError(
+ "`cross_attention_gate` is required for Idefics cross attention module to zero-out the cross-attention hidden_states attending to no images."
+ )
+
+ if past_key_values is not None:
+ raise NotImplementedError("Past key value states are not implemented for Idefics cross attention module.")
+
+ residual = hidden_states
+
+ hidden_states = self.input_layernorm(hidden_states)
+
+ # Self Attention
+ hidden_states, _ = self.cross_attn(
+ hidden_states=hidden_states,
+ key_value_states=image_hidden_states,
+ attention_mask=image_attention_mask,
+ **kwargs,
+ )
+ hidden_states = nn.functional.dropout(hidden_states, p=self.config, training=self.training)
+ # Fill in zeros for cross_attention hidden_states of tokens attending to no images
+ hidden_states = hidden_states.masked_fill((cross_attention_gate == 0)[:, :, None], 0.0)
+ hidden_states = residual + self.act_cross_attn(self.alpha_cross_attn) * hidden_states
+
+ # Fully Connected
+ residual = hidden_states
+ hidden_states = self.post_attention_layernorm(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = nn.functional.dropout(hidden_states, p=self.config, training=self.training)
+ hidden_states = residual + self.act_dense(self.alpha_dense) * hidden_states
+
+ return hidden_states
+
+
+@auto_docstring
+class IdeficsPreTrainedModel(PreTrainedModel):
+ config: IdeficsConfig
+ base_model_prefix = "model"
+ input_modalities = ("image", "text")
+ supports_gradient_checkpointing = True
+ _no_split_modules = ["IdeficsDecoderLayer", "IdeficsGatedCrossAttentionLayer"]
+ _supports_sdpa = True
+
+ _supports_flash_attn = False # only eager/sdpa creation is supported
+ _can_compile_fullgraph = False # IDEFICS cannot compile due to dynamic control flow when checking inputs
+ _supports_attention_backend = True
+
+ _can_record_outputs = {
+ "hidden_states": IdeficsDecoderLayer,
+ "attentions": OutputRecorder(IdeficsAttention, index=1, layer_name="self_attn"),
+ }
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ # important: this ported version of Idefics 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.
+ super()._init_weights(module)
+ if isinstance(module, IdeficsVisionEmbeddings):
+ init.normal_(module.class_embedding)
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+ elif isinstance(module, IdeficsGatedCrossAttentionLayer):
+ if self.config.alpha_initializer == "zeros":
+ init.zeros_(module.alpha_cross_attn)
+ init.zeros_(module.alpha_dense)
+ elif self.config.alpha_initializer == "ones":
+ init.ones_(module.alpha_cross_attn)
+ init.ones_(module.alpha_dense)
+ elif self.config.alpha_initializer in {"normal", "gaussian", "random"}:
+ init.normal_(module.alpha_cross_attn, mean=0.0, std=self.config.alphas_initializer_range)
+ init.normal_(module.alpha_dense, mean=0.0, std=self.config.alphas_initializer_range)
+ elif isinstance(module, IdeficsPerceiverResampler):
+ init.normal_(module.latents)
+ elif isinstance(module, IdeficsEmbedding):
+ inv_freq = 1.0 / (module.base ** (torch.arange(0, module.dim, 2) / module.dim))
+ init.copy_(module.inv_freq, inv_freq)
+ t = torch.arange(module.max_position_embeddings).type_as(inv_freq)
+ freqs = torch.einsum("i,j->ij", t, 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)
+ init.copy_(module.cos_cached, emb.cos())
+ init.copy_(module.sin_cached, emb.sin())
+
+
+@auto_docstring
+class IdeficsModel(IdeficsPreTrainedModel):
+ """
+ Transformer decoder consisting of `config.num_hidden_layers` layers. Each layer is a [`IdeficsDecoderLayer`]
+
+ Args:
+ config: IdeficsConfig
+ """
+
+ def __init__(self, config: IdeficsConfig):
+ super().__init__(config)
+ self.config = config
+ self.padding_idx = config.pad_token_id
+ self.vocab_size = config.vocab_size
+
+ self.embed_tokens = IdeficsDecoupledEmbedding(
+ 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,
+ )
+
+ self.image_size = config.vision_config.image_size
+ self.vision_config = config.vision_config
+ # The module using it is not a PreTrainedModel subclass so we need this
+ self.vision_config._attn_implementation = config._attn_implementation
+ self.vision_model = IdeficsVisionTransformer(config.vision_config)
+
+ # Perceiver Resampler
+ if config.use_resampler:
+ perceiver_config = config.perceiver_config
+ self.perceiver_resampler = IdeficsPerceiverResampler(
+ config,
+ config.vision_config.embed_dim,
+ perceiver_config.resampler_depth,
+ perceiver_config.resampler_n_heads,
+ perceiver_config.resampler_head_dim,
+ perceiver_config.resampler_n_latents,
+ )
+
+ self.layers = nn.ModuleList(
+ [IdeficsDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)]
+ )
+
+ self.cross_layer_interval = config.cross_layer_interval
+ num_cross_layers = config.num_hidden_layers // self.cross_layer_interval
+ self.gated_cross_attn_layers = nn.ModuleList(
+ [IdeficsGatedCrossAttentionLayer(config, layer_idx=i) for i in range(num_cross_layers)]
+ )
+ self.gradient_checkpointing = False
+
+ self.norm = IdeficsRMSNorm(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 freeze_vision_layers(self, module_exceptions=()):
+ freeze_model(self.vision_model, module_exceptions=module_exceptions)
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ image_encoder_embeddings: torch.FloatTensor | None = None,
+ perceiver_embeddings: torch.FloatTensor | None = None,
+ image_attention_mask: torch.Tensor | None = None,
+ use_cache: bool | None = None,
+ interpolate_pos_encoding: bool | None = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | IdeficsBaseModelOutputWithPast:
+ r"""
+ image_encoder_embeddings (`torch.FloatTensor`, *optional*):
+ The output of the image encoder.
+ perceiver_embeddings (`torch.FloatTensor`, *optional*):
+ The output of the perceiver resampler.
+ image_attention_mask (`torch.LongTensor`, *optional*):
+ The attention mask for the image encoder.
+ """
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+ if (input_ids is None) ^ (inputs_embeds is not None):
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
+
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_tokens(input_ids)
+
+ if use_cache and past_key_values is None:
+ past_key_values = DynamicCache(config=self.config)
+
+ seq_length = inputs_embeds.shape[1]
+ past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0
+ seq_length_with_past = seq_length + past_key_values_length
+
+ 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)
+ position_ids = position_ids[:, -seq_length:]
+ elif position_ids is None:
+ position_ids = torch.arange(seq_length, device=inputs_embeds.device) + past_key_values_length
+ position_ids = position_ids.unsqueeze(0)
+
+ if sum(x is None for x in [pixel_values, image_encoder_embeddings, perceiver_embeddings]) != 2:
+ raise ValueError(
+ "Exactly 1 of pixel_values, image_encoder_embeddings or perceiver_embeddings has to be not-None."
+ )
+
+ elif pixel_values is not None:
+ pixel_values = pixel_values.to(dtype=self.dtype, device=device) # fp16 compatibility
+ batch_size, num_images = pixel_values.shape[:2]
+ pixel_values = pixel_values.contiguous().view(batch_size * num_images, *pixel_values.shape[2:])
+
+ # Get sequence from the vision encoder
+ image_hidden_states = self.vision_model(
+ pixel_values=pixel_values, interpolate_pos_encoding=interpolate_pos_encoding
+ ).last_hidden_state
+
+ elif image_encoder_embeddings is not None:
+ batch_size, num_images, image_seq_len, image_hidden_size = image_encoder_embeddings.size()
+ image_hidden_states = image_encoder_embeddings.to(dtype=self.dtype, device=device)
+ image_hidden_states = image_hidden_states.view(batch_size * num_images, image_seq_len, image_hidden_size)
+
+ if self.config.use_resampler:
+ if perceiver_embeddings is None:
+ perceiver_embeddings = self.perceiver_resampler(image_hidden_states)
+ image_seq_len, image_hidden_size = perceiver_embeddings.size(1), perceiver_embeddings.size(2)
+ else:
+ batch_size, num_images, image_seq_len, image_hidden_size = perceiver_embeddings.size()
+ image_hidden_states = perceiver_embeddings
+ elif perceiver_embeddings is None:
+ image_seq_len, image_hidden_size = image_hidden_states.size(1), image_hidden_states.size(2)
+ else:
+ raise ValueError("If `perceiver_embeddings` are passed, use_resampler should be True")
+
+ image_hidden_states = image_hidden_states.view(batch_size, num_images * image_seq_len, image_hidden_size)
+ # # Hack to use the model in full language modeling mode
+ # image_attention_mask = torch.zeros(batch_size, seq_length, 1, dtype=torch.long, device=image_hidden_states.device)
+ # Make image_attention_mask compatible with hidden states
+ text_seq_len = image_attention_mask.size(1)
+ image_attention_mask = image_attention_mask.unsqueeze(-1)
+ image_attention_mask = image_attention_mask.repeat(1, 1, 1, image_seq_len)
+ image_attention_mask = image_attention_mask.view(batch_size, text_seq_len, num_images * image_seq_len)
+
+ if image_hidden_states is not None:
+ image_batch_size, image_sequence_length, _ = image_hidden_states.size()
+ image_hidden_shape = (image_batch_size, image_sequence_length)
+ if image_attention_mask is None:
+ image_attention_mask = torch.ones(image_hidden_shape, device=device)
+ image_attention_mask = self.invert_attention_mask(image_attention_mask)
+ else:
+ image_attention_mask = None
+
+ # cross_attention_gate:
+ # For any tokens attending to no images, the hidden_states coming out of the cross-attention should be zeroed-out.
+ # `image_attention_mask` has shape [bsz, 1, num_images, hidden_size] with elements equal to either 0.0 or a very negative number.
+ # If any of the elements are 0.0, then the token is attending to at least one image and the gate value is 1. Otherwise the gate value is 0.
+ # `cross_attention_gate` has shape [bsz, seq_len] with elements equal to either 0.0 or 1.0.
+ cross_attention_gate = ((((image_attention_mask == 0.0).any(dim=-1)).to(dtype=self.dtype)).squeeze(dim=1)).to(
+ device
+ )
+
+ # embed positions
+ if attention_mask is None:
+ attention_mask = torch.ones(
+ (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
+ )
+
+ causal_mask = create_causal_mask(
+ config=self.config,
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ past_key_values=past_key_values,
+ position_ids=position_ids,
+ )
+
+ hidden_states = inputs_embeds
+
+ for idx, decoder_layer in enumerate(self.layers):
+ # TODO(ls): Add cross attention values to respective lists
+ if idx % self.cross_layer_interval == 0:
+ cross_attn_block = self.gated_cross_attn_layers[idx // self.cross_layer_interval]
+ hidden_states = cross_attn_block(
+ hidden_states,
+ causal_mask,
+ image_hidden_states,
+ image_attention_mask=image_attention_mask,
+ cross_attention_gate=cross_attention_gate,
+ past_key_values=None, # not implemented
+ **kwargs,
+ )
+
+ hidden_states = decoder_layer(
+ hidden_states,
+ attention_mask=causal_mask,
+ position_ids=position_ids,
+ past_key_values=past_key_values,
+ **kwargs,
+ )
+
+ hidden_states = self.norm(hidden_states)
+ image_hidden_states = image_hidden_states.view(batch_size, num_images, image_seq_len, image_hidden_size)
+
+ return IdeficsBaseModelOutputWithPast(
+ last_hidden_state=hidden_states,
+ image_hidden_states=image_hidden_states,
+ past_key_values=past_key_values,
+ )
+
+
+class IdeficsForVisionText2Text(IdeficsPreTrainedModel, GenerationMixin):
+ _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+ def __init__(self, config, vision_model=None):
+ super().__init__(config)
+ self.model = IdeficsModel(config)
+
+ self.lm_head = IdeficsDecoupledLinear(
+ 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,
+ )
+ if config.additional_vocab_size > 0:
+ self._tied_weights_keys = {
+ "lm_head.weight": "model.embed_tokens.weight",
+ "lm_head.additional_fc.weight": "model.embed_tokens.additional_embedding.weight",
+ }
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.Tensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ past_key_values: Cache | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ pixel_values: torch.FloatTensor | None = None,
+ image_encoder_embeddings: torch.FloatTensor | None = None,
+ perceiver_embeddings: torch.FloatTensor | None = None,
+ image_attention_mask: torch.Tensor | None = None,
+ labels: torch.LongTensor | None = None,
+ use_cache: bool | None = None,
+ interpolate_pos_encoding: bool | None = False,
+ logits_to_keep: int | torch.Tensor = 0,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | IdeficsCausalLMOutputWithPast:
+ r"""
+ image_encoder_embeddings (`torch.FloatTensor`, *optional*):
+ The output of the image encoder.
+ perceiver_embeddings (`torch.FloatTensor`, *optional*):
+ The output of the perceiver resampler.
+ image_attention_mask (`torch.LongTensor`, *optional*):
+ The attention mask for the image encoder.
+ 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]`.
+
+ Example:
+
+ ```python
+ >>> from transformers import AutoProcessor, IdeficsForVisionText2Text
+
+ >>> model = IdeficsForVisionText2Text.from_pretrained("HuggingFaceM4/idefics-9b")
+ >>> processor = AutoProcessor.from_pretrained("HuggingFaceM4/idefics-9b")
+
+ >>> dogs_image_url_1 = "https://huggingface.co/datasets/hf-internal-testing/fixtures_nlvr2/raw/main/image1.jpeg"
+ >>> dogs_image_url_2 = "https://huggingface.co/datasets/hf-internal-testing/fixtures_nlvr2/raw/main/image2.jpeg"
+
+ >>> prompts = [
+ ... [
+ ... "User:",
+ ... dogs_image_url_1,
+ ... "Describe this image.\nAssistant: An image of two dogs.\n",
+ ... "User:",
+ ... dogs_image_url_2,
+ ... "Describe this image.\nAssistant:",
+ ... ]
+ ... ]
+ >>> inputs = processor(prompts, return_tensors="pt")
+ >>> generate_ids = model.generate(**inputs, max_new_tokens=6)
+ >>> processor.batch_decode(generate_ids, skip_special_tokens=True)
+ ```"""
+ outputs: IdeficsBaseModelOutputWithPast = 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_encoder_embeddings=image_encoder_embeddings,
+ perceiver_embeddings=perceiver_embeddings,
+ image_attention_mask=image_attention_mask,
+ use_cache=use_cache,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ return_dict=True,
+ **kwargs,
+ )
+
+ hidden_states = outputs.last_hidden_state
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
+
+ return IdeficsCausalLMOutputWithPast(
+ 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,
+ attention_mask=None,
+ position_ids=None,
+ inputs_embeds=None,
+ past_key_values=None,
+ pixel_values=None,
+ image_hidden_states=None,
+ image_attention_mask=None,
+ use_cache=None,
+ **kwargs,
+ ):
+ # Overwritten -- custom processing based on `config.use_resampler`
+
+ images_kwargs = {}
+ if image_hidden_states is not None:
+ if self.config.use_resampler:
+ images_kwargs["perceiver_embeddings"] = image_hidden_states
+ else:
+ images_kwargs["image_encoder_embeddings"] = image_hidden_states
+ else:
+ images_kwargs["pixel_values"] = pixel_values
+ images_kwargs["interpolate_pos_encoding"] = kwargs.pop("interpolate_pos_encoding", False)
+
+ model_inputs = super().prepare_inputs_for_generation(
+ input_ids,
+ past_key_values=past_key_values,
+ attention_mask=attention_mask,
+ inputs_embeds=inputs_embeds,
+ position_ids=position_ids,
+ use_cache=use_cache,
+ image_attention_mask=image_attention_mask,
+ **images_kwargs,
+ **kwargs,
+ )
+
+ if image_attention_mask is not None and inputs_embeds is None:
+ seq_length = model_inputs["input_ids"].shape[1]
+ model_inputs["image_attention_mask"] = image_attention_mask[:, -seq_length:]
+
+ return model_inputs
+
+ def _update_model_kwargs_for_generation(
+ self,
+ outputs: ModelOutput,
+ model_kwargs: dict[str, Any],
+ is_encoder_decoder: bool = False,
+ **kwargs,
+ ) -> dict[str, Any]:
+ model_kwargs = super()._update_model_kwargs_for_generation(
+ outputs,
+ model_kwargs,
+ is_encoder_decoder,
+ **kwargs,
+ )
+
+ if "image_attention_mask" in model_kwargs:
+ image_attention_mask = model_kwargs["image_attention_mask"]
+ last_mask = image_attention_mask[:, -1, :].unsqueeze(1)
+ if model_kwargs.get("use_cache", True):
+ model_kwargs["image_attention_mask"] = last_mask
+ else:
+ model_kwargs["image_attention_mask"] = torch.cat([image_attention_mask, last_mask], dim=1)
+
+ # Get the precomputed image_hidden_states
+ model_kwargs["image_hidden_states"] = outputs.image_hidden_states
+ return model_kwargs
+
+
+__all__ = ["IdeficsForVisionText2Text", "IdeficsModel", "IdeficsPreTrainedModel"]
diff --git a/third_party/transformers/src/transformers/models/idefics/perceiver.py b/third_party/transformers/src/transformers/models/idefics/perceiver.py
new file mode 100644
index 0000000000000000000000000000000000000000..96b66ae0b3daca8ec0d9acd0be2dda615c332e13
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/idefics/perceiver.py
@@ -0,0 +1,187 @@
+# This code was adapted from https://github.com/lucidrains/flamingo-pytorch licensed under the MIT License.
+#
+# MIT License
+#
+# Copyright (c) 2020 The Google AI Language Team Authors, The HuggingFace Inc. team and github/lonePatient
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+
+"""
+
+Generic interface to various configurations of the Perceiver Resampler, that simply takes in a series of (potentially
+time-indexed) contextual embeddings, and "resamples" (compresses) them down to a pre-specified number of latents! Note
+that the Perceiver in general resamples based solely off the *long-range* context; there's a nice opportunity here to
+prime the Perceiver Resampler with say a single layer's worth of language embeddings (the target domain), and use that
+to softly "retrieve & compress" what we need --> this would be a novel contribution we should explore.
+
+References:
+ - DeepMind's Flamingo: https://www.deepmind.com/blog/tackling-multiple-tasks-with-a-single-visual-language-model
+ - Code borrowed w/ love from: https://github.com/lucidrains/flamingo-pytorch
+
+"""
+
+import torch
+import torch.nn as nn
+
+from .configuration_idefics import IdeficsConfig
+
+
+class IdeficsPerceiverResampler(nn.Module):
+ def __init__(
+ self, config: IdeficsConfig, embed_dim: int, depth: int, n_heads: int, head_dim: int, n_latents: int
+ ) -> 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.
+
+ Args:
+ config (`IdeficsConfig`): config object
+ embed_dim (`int`): The size of each embedding vector
+ depth (`int`): Depth of the Perceiver Resampler (Transformer w/ cross attention). Should be shallow (< 3).
+ n_heads (`int`): Number of heads in each Transformer block (for multi-headed self-attention).
+ head_dim (`int`): Dimensionality of each head projection in the Transformer block.
+ n_latents (`int`):
+ 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 = config.perceiver_config.qk_layer_norms_perceiver
+
+ # Create Latents for Perceiver
+ self.latents = nn.Parameter(torch.randn(self.n_latents, self.embed_dim), requires_grad=True)
+
+ self.intermediate_dim = (
+ self.embed_dim * 4
+ if not hasattr(config.vision_config, "embed_dim")
+ else config.vision_config.embed_dim * 4
+ )
+ # Create Transformer Blocks
+ self.blocks = nn.ModuleList(
+ [
+ nn.ModuleList(
+ [
+ IdeficsPerceiverAttention(self.embed_dim, self.n_heads, self.head_dim, self.qk_layer_norms),
+ IdeficsMLP(self.intermediate_dim, config),
+ ]
+ )
+ 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"""
+ # einsum.repeat(self.latents, "seq embed -> bsz seq embed", bsz=context.shape[0])
+ latents = self.latents.repeat(context.shape[0], 1, 1)
+
+ # 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 IdeficsPerceiverAttention(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, 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!
+
+ Args:
+ context (`torch.Tensor`):
+ Tensor of shape `[bsz, seq, embed_dim]` representing long-form context to resample.
+ latents (`torch.Tensor`):
+ Tensor of shape `[bsz, n_latents, embed_dim]` representing fixed length latents to compress to.
+
+ Returns:
+ `torch.Tensor`: 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)
+ batch_size, seq_length, embed_dim = context.shape[:3]
+
+ # 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)]
+ # einsum.rearrange(x, "bsz seq (heads embed) -> bsz heads seq embed", heads=self.n_heads)
+ q, k, v = [x.reshape(batch_size, x.shape[1], self.n_heads, self.head_dim).transpose(1, 2) 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)
+ # einsum.rearrange(resampled, "bsz heads seq embed -> bsz seq (heads embed)", heads=self.n_heads)
+ return self.output_proj(resampled.transpose(1, 2).flatten(-2))
+
+
+class IdeficsMLP(nn.Module):
+ def __init__(self, intermediate_size, config: IdeficsConfig):
+ """Simple MLP block with intermediate_size and embedding size"""
+ super().__init__()
+ self.embed_dim = config.vision_config.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: tuple[torch.FloatTensor] | None) -> 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
diff --git a/third_party/transformers/src/transformers/models/idefics/processing_idefics.py b/third_party/transformers/src/transformers/models/idefics/processing_idefics.py
new file mode 100644
index 0000000000000000000000000000000000000000..b774d10b35c7e2f86d940ba723e1522244778f05
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/idefics/processing_idefics.py
@@ -0,0 +1,416 @@
+# Copyright 2022 The HuggingFace Inc. team.
+#
+# 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.
+"""
+Processor class for IDEFICS.
+"""
+
+from urllib.parse import urlparse
+
+from ...feature_extraction_utils import BatchFeature
+from ...image_utils import ImageInput
+from ...processing_utils import (
+ ProcessingKwargs,
+ ProcessorMixin,
+ TextKwargs,
+ Unpack,
+)
+from ...tokenization_utils_base import PreTokenizedInput, TextInput
+from ...utils import auto_docstring, is_torch_available
+
+
+if is_torch_available():
+ import torch
+
+
+IMAGE_TOKEN = ""
+
+
+class IdeficsTextKwargs(TextKwargs, total=False):
+ """
+ add_eos_token (`bool`, *optional*, defaults to `False`):
+ Whether to add an end-of-sequence token at the end of the text input. When enabled, an EOS token is
+ appended to mark the end of the text sequence, which is useful for generation tasks.
+ add_end_of_utterance_token (`bool`, *optional*):
+ Whether to add an end-of-utterance token to mark the end of a user's message in conversational contexts.
+ This token helps the model distinguish between different utterances in a multi-turn conversation and is
+ particularly important for chat-based models.
+ """
+
+ add_eos_token: bool | None
+ add_end_of_utterance_token: bool | None
+
+
+class IdeficsProcessorKwargs(ProcessingKwargs, total=False):
+ text_kwargs: IdeficsTextKwargs
+ _defaults = {
+ "text_kwargs": {
+ "add_special_tokens": False,
+ "padding": "longest",
+ "add_eos_token": False,
+ },
+ "common_kwargs": {"return_tensors": "pt"},
+ }
+
+
+# copied from m4.training.packing
+def incremental_to_binary_attention_mask(incremental_mask, return_tensors, num_classes=-1):
+ # Set elements >= num_classes to -1
+ if num_classes != -1:
+ if return_tensors == "pt":
+ incremental_mask[incremental_mask >= num_classes] = -1
+
+ # Create mask for negative values
+ if return_tensors == "pt":
+ negatives = incremental_mask == -1
+ incremental_mask[negatives] = 0
+ attn_mask = torch.nn.functional.one_hot(incremental_mask, num_classes=num_classes)
+ attn_mask[negatives, :] = 0
+
+ return attn_mask
+
+
+# copied from m4.training.packing
+def image_attention_mask_for_packed_input_ids(input_ids, tokenizer, return_tensors):
+ if return_tensors == "pt":
+ return image_attention_mask_for_packed_input_ids_pt(input_ids, tokenizer)
+
+
+def image_attention_mask_for_packed_input_ids_pt(input_ids, tokenizer):
+ image_attention_mask = torch.full_like(input_ids, fill_value=-1)
+ next_image_attention_mask = torch.full_like(input_ids, fill_value=-1)
+ image_token_id = tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
+ eod_token_id = tokenizer.eos_token_id
+ for batch_idx in range(input_ids.size(0)):
+ count = -1
+ seen_eod = False
+ for idx, token_id in enumerate(input_ids[batch_idx]):
+ if token_id == image_token_id:
+ count += 1
+ image_attention_mask[batch_idx][idx] = count
+ seen_eod = False
+ else:
+ image_attention_mask[batch_idx][idx] = count
+
+ if seen_eod:
+ image_attention_mask[batch_idx][idx] = -1
+
+ if token_id == eod_token_id:
+ seen_eod = True
+
+ for batch_idx in range(input_ids.size(0)):
+ count = -1
+ seen_eod = False
+ for idx in range(input_ids[batch_idx].size(0) - 1, -1, -1):
+ token_id = input_ids[batch_idx][idx]
+ if token_id == image_token_id:
+ count += 1
+ next_image_attention_mask[batch_idx][idx] = count
+ seen_eod = False
+ else:
+ next_image_attention_mask[batch_idx][idx] = count
+
+ if token_id == eod_token_id:
+ seen_eod = True
+
+ if seen_eod:
+ next_image_attention_mask[batch_idx][idx] = -1
+
+ non_negative_indices = next_image_attention_mask[batch_idx] != -1
+ next_image_attention_mask[batch_idx][non_negative_indices] -= count
+ next_image_attention_mask[batch_idx][non_negative_indices] *= -1
+
+ return image_attention_mask, next_image_attention_mask
+
+
+def is_url(string):
+ """Checks if the passed string contains a valid url and nothing else. e.g. if space is included it's immediately
+ invalidated the url"""
+ if " " in string:
+ return False
+ result = urlparse(string)
+ return all([result.scheme, result.netloc])
+
+
+@auto_docstring
+class IdeficsProcessor(ProcessorMixin):
+ def __init__(self, image_processor, tokenizer=None, image_size=224, add_end_of_utterance_token=None, **kwargs):
+ r"""
+ image_size (int, *optional*, defaults to 224):
+ The size of the image to be processed.
+ add_end_of_utterance_token (bool, *optional*, defaults to None):
+ Whether to add the end of utterance token to the text.
+ """
+ super().__init__(image_processor, tokenizer)
+ self.image_token_id = (
+ tokenizer.image_token_id
+ if hasattr(tokenizer, "image_token")
+ else tokenizer.convert_tokens_to_ids(IMAGE_TOKEN)
+ )
+
+ self.default_image_dims = (
+ self.image_processor.image_num_channels,
+ self.image_processor.image_size,
+ self.image_processor.image_size,
+ )
+
+ self.tokenizer_was_trained_with_end_of_utterance_token = (
+ "" in self.tokenizer.special_tokens_map.get("additional_special_tokens", [])
+ )
+
+ @auto_docstring
+ def __call__(
+ self,
+ images: ImageInput | list[ImageInput] | str | list[str] | list[list[str]] = None,
+ text: TextInput
+ | PreTokenizedInput
+ | list[TextInput]
+ | list[PreTokenizedInput]
+ | list[list[TextInput]]
+ | list[list[PreTokenizedInput]] = None,
+ **kwargs: Unpack[IdeficsProcessorKwargs],
+ ) -> BatchFeature:
+ r"""
+ Returns:
+ a dict with entries: `input_ids`, `attention_mask`, `pixel_values`, `image_attention_mask` which can be
+ directly passed to `model.generate`
+
+ Detailed explanation:
+
+ Each entry in `text` is either a text to be passed as is or an image that will be processed.
+
+ An image can be either an image object (`PIL.Image`) or a url from which the image can be retrieved.
+
+ When the processor encounters an image it'll inject ``
+ entry into the prompt.
+
+ Example:
+
+ ```python
+ checkpoint = "HuggingFaceM4/idefics-9b"
+ processor = AutoProcessor.from_pretrained(checkpoint)
+ url = "https://hips.hearstapps.com/hmg-prod/images/cute-photos-of-cats-in-grass-1593184777.jpg"
+ img = processor.image_processor.fetch_images([url])[0]
+
+ prompts = [
+ "User:",
+ img,
+ "Describe this image.\nAssistant: An image of two kittens in grass.\n",
+ "User:",
+ "https://hips.hearstapps.com/hmg-prod/images/dog-puns-1581708208.jpg",
+ "Describe this image.\nAssistant:",
+ ]
+
+ inputs = processor(text=prompts, return_tensors="pt")
+ generated_ids = model.generate(**inputs, max_length=100)
+ generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+ ```
+
+ In this example the `prompts` will be converted into:
+
+ ```
+ User:Describe this image.
+ Assistant: An image of two kittens in grass.
+ User:Describe this image.
+ Assistant:'
+ ```
+
+ and the two images will be massaged using [`IdeficsImageProcessor.__call__`] method and placed inside the
+ `pixel_values` dict entry of the return value.
+
+ This example also exemplifies that images can be passed as objects or as text urls. It can be seen that the
+ first image is passed as object and the second one as a url.
+
+ To do training do:
+
+ ```python
+ image_transform = transforms.Compose(
+ [
+ transforms.RandomResizedCrop(
+ (w, h), scale=(0.9, 1.0), interpolation=transforms.InterpolationMode.BICUBIC
+ ),
+ transforms.ToTensor(),
+ transforms.Normalize(mean=self.image_mean, std=self.image_std),
+ ]
+ )
+ inputs = processor(text=prompts, transform=image_transform, return_tensors="pt")
+ ```
+
+ In order to help debug prompt generation enable `debug=True` which will show you what's happening.
+
+ """
+ if images is None and text is None:
+ raise ValueError("You need to specify either `text` or `images` and `text`.")
+
+ if images is None:
+ # assuming the user wants to use the old behavior with prompts as the only argument
+ prompts = text
+ elif text is not None:
+ # Assuming image-text-to-text behavior:
+ # Check if batched images are provided
+ if not isinstance(images, (list, tuple)):
+ images = [images]
+ if isinstance(text, str):
+ text = [text]
+ # Check if batched images and text are in the correct format
+ if isinstance(text, (list, tuple)) and len(text) != len(images):
+ raise ValueError(
+ "When providing both images and text arguments, the number of text prompts should be the same as the number of images."
+ "If you want to have several images per prompt, images should be nested as such: images=[[img1, img2], [img3, img4], ...] for text=[prompt1, prompt2, ...]."
+ )
+ # Check that only text is present in the prompts
+ if not all(isinstance(i, str) for i in text):
+ raise ValueError("When using the image-text-to-text behavior, the prompts should only contain text.")
+ if isinstance(images[0], (list, tuple)):
+ # if nested images, un-nest each sublist and create `prompts`
+ prompts = [[sample, *image_list] for image_list, sample in zip(images, text)]
+ else:
+ prompts = list(zip(images, text))
+
+ output_kwargs = self._merge_kwargs(
+ IdeficsProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+
+ add_eos_token = output_kwargs["text_kwargs"].pop("add_eos_token", False)
+ add_end_of_utterance_token = output_kwargs["text_kwargs"].pop("add_end_of_utterance_token", None)
+
+ # if the value isn't overridden by the user, check if the tokenizer was trained with this token and then use it
+ if add_end_of_utterance_token is None:
+ add_end_of_utterance_token = self.tokenizer_was_trained_with_end_of_utterance_token
+ # turn non-batched prompts into batched
+ if not any(isinstance(i, (list, tuple)) for i in prompts):
+ prompts = [prompts]
+
+ fake_token = ""
+ image_token = ""
+ end_of_utterance_token = ""
+
+ def image_tokens(last_was_image):
+ if last_was_image:
+ return image_token + fake_token
+ else:
+ return fake_token + image_token + fake_token
+
+ all_prompts = []
+ all_images = []
+ for sample in prompts:
+ # the model was trained on samples starting with
+ full_text = f"{self.tokenizer.bos_token}"
+
+ # an image can either be an image object in the item or the url, everything else is a verbatim prompt text
+ image_objects = []
+ last_was_image = False
+ last_was_text = False
+ for i, item in enumerate(sample):
+ if i > 0:
+ last_was_text = bool(not last_was_image)
+
+ if isinstance(item, str):
+ item = item.strip(" ")
+ if is_url(item):
+ image = self.image_processor.fetch_images(item)
+ full_text += image_tokens(last_was_image)
+ image_objects.append(image)
+ last_was_image = True
+ else:
+ # we add end_of_utterance_token between each subsequent text prompts (but not at the last one!)
+ if add_end_of_utterance_token and last_was_text:
+ full_text += end_of_utterance_token
+ full_text += item
+ last_was_image = False
+ else:
+ # must be an image obj
+ full_text += image_tokens(last_was_image)
+ image_objects.append(item)
+ last_was_image = True
+
+ if add_eos_token:
+ full_text += self.tokenizer.eos_token
+
+ if len(image_objects) > 0:
+ image_objects = self.image_processor(image_objects, **output_kwargs["images_kwargs"])
+
+ all_prompts.append(full_text)
+ all_images.append(image_objects)
+
+ # For BC
+ return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", "pt")
+ text_encoding = self.tokenizer(all_prompts, **output_kwargs["text_kwargs"])
+ all_texts = text_encoding["input_ids"]
+ all_attention_masks = text_encoding["attention_mask"]
+
+ # max_num_images has to be at least 1 even when there are no images
+ max_num_images = max(len(x) for x in all_images)
+ max_num_images = max(1, max_num_images)
+
+ at_least_one_image = sum(len(x) for x in all_images) > 0
+ output_input_ids = []
+ output_images = []
+ output_attention_masks = []
+
+ for text_single, attention_mask, extracted_images in zip(all_texts, all_attention_masks, all_images):
+ padded_input_ids = text_single
+ image_count = padded_input_ids.count(self.image_token_id)
+ local_max_num_images = min(image_count, max_num_images)
+ current_images = extracted_images[:local_max_num_images]
+
+ if len(current_images) > 0:
+ if return_tensors == "pt":
+ padded_image_tensor = torch.zeros(max_num_images, *current_images.size()[1:])
+ padded_image_tensor[: current_images.size(0)] = current_images
+ else:
+ if return_tensors == "pt":
+ padded_image_tensor = torch.zeros(max_num_images, *self.default_image_dims)
+
+ output_images.append(padded_image_tensor)
+ if return_tensors == "pt":
+ output_input_ids.append(torch.tensor(padded_input_ids))
+ output_attention_masks.append(torch.tensor(attention_mask))
+
+ if return_tensors == "pt":
+ output_input_ids = torch.stack(output_input_ids)
+ output_images = torch.stack(output_images)
+ output_attention_masks = torch.stack(output_attention_masks)
+
+ if at_least_one_image:
+ image_attention_mask, _ = image_attention_mask_for_packed_input_ids(
+ output_input_ids, self.tokenizer, return_tensors
+ )
+ image_attention_mask = incremental_to_binary_attention_mask(
+ image_attention_mask, return_tensors, num_classes=max_num_images
+ )
+ else:
+ # in full language mode we set the image mask to all-0s
+ if return_tensors == "pt":
+ image_attention_mask = torch.zeros(
+ output_input_ids.shape[0], output_input_ids.shape[1], 1, dtype=torch.bool
+ )
+ return BatchFeature(
+ data={
+ "input_ids": output_input_ids,
+ "attention_mask": output_attention_masks,
+ "pixel_values": output_images,
+ "image_attention_mask": image_attention_mask,
+ }
+ )
+
+ @property
+ def model_input_names(self):
+ tokenizer_input_names = self.tokenizer.model_input_names
+ image_processor_input_names = self.image_processor.model_input_names
+ return list(tokenizer_input_names + image_processor_input_names + ["image_attention_mask"])
+
+
+__all__ = ["IdeficsProcessor"]
diff --git a/third_party/transformers/src/transformers/models/idefics/vision.py b/third_party/transformers/src/transformers/models/idefics/vision.py
new file mode 100644
index 0000000000000000000000000000000000000000..dbd22df3d7c2d409875146a69ff13f3beaa17f59
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/idefics/vision.py
@@ -0,0 +1,395 @@
+# Copyright 2021 The OpenAI Team Authors and The HuggingFace Team. All rights reserved.
+#
+# 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 IdeficsVision model: a copy of CLIPVisionModel using a simpler config object"""
+
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from ...activations import ACT2FN
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS
+from ...processing_utils import Unpack
+from ...utils import (
+ ModelOutput,
+ TransformersKwargs,
+ logging,
+)
+from .configuration_idefics import IdeficsVisionConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+class IdeficsVisionModelOutput(ModelOutput):
+ """
+ Base class for vision model's outputs that also contains image embeddings of the pooling of the last hidden states.
+
+ Args:
+ image_embeds (`torch.FloatTensor` of shape `(batch_size, output_dim)` *optional* returned when model is initialized with `with_projection=True`):
+ The image embeddings obtained by applying the projection layer to the pooler_output.
+ 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.
+ 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_embeds: torch.FloatTensor | None = None
+ last_hidden_state: torch.FloatTensor | None = None
+ hidden_states: tuple[torch.FloatTensor, ...] | None = None
+ attentions: tuple[torch.FloatTensor, ...] | None = None
+
+
+# Adapted from transformers.models.clip.modeling_clip.CLIPVisionEmbeddings
+class IdeficsVisionEmbeddings(nn.Module):
+ def __init__(self, config: IdeficsVisionConfig):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.image_size = config.image_size
+ self.patch_size = config.patch_size
+
+ self.class_embedding = nn.Parameter(torch.randn(self.embed_dim))
+
+ self.patch_embedding = nn.Conv2d(
+ in_channels=config.num_channels,
+ out_channels=self.embed_dim,
+ kernel_size=self.patch_size,
+ stride=self.patch_size,
+ bias=False,
+ )
+
+ self.num_patches = (self.image_size // self.patch_size) ** 2
+ self.num_positions = self.num_patches + 1
+ self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim)
+ self.register_buffer("position_ids", torch.arange(self.num_positions).expand((1, -1)), persistent=False)
+
+ # Heavily inspired from https://github.com/huggingface/transformers/blob/v4.33.0/src/transformers/models/vit/modeling_vit.py#L82
+ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+ """
+ This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
+ resolution images.
+
+ Source:
+ https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
+ """
+
+ num_patches = embeddings.shape[1] - 1
+ pos_embed = self.position_embedding(self.position_ids)
+ num_positions = pos_embed.shape[1] - 1
+ if num_patches == num_positions and height == width:
+ return pos_embed
+ class_pos_embed = pos_embed[:, 0]
+ patch_pos_embed = pos_embed[:, 1:]
+
+ embed_dim = embeddings.shape[-1]
+ num_h_patches = height // self.config.patch_size
+ num_w_patches = width // self.config.patch_size
+ # we add a small number to avoid floating point error in the interpolation
+ # see discussion at https://github.com/facebookresearch/dino/issues/8
+ num_h_patches, num_w_patches = num_h_patches + 0.1, num_w_patches + 0.1
+ sqrt_num_positions = math.sqrt(num_positions)
+ patch_pos_embed = patch_pos_embed.reshape(1, int(sqrt_num_positions), int(sqrt_num_positions), embed_dim)
+ patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+ fp32_upcasting = patch_pos_embed.dtype == torch.bfloat16
+ if fp32_upcasting:
+ logger.warning_once(
+ "Upcasting patch_pos_embed to fp32 for interpolation since `upsample_bicubic2d_out_frame` in nn.functional.interpolate "
+ "is not implemented for 'torch.bfloat16' dtype. This will result in a slight overhead."
+ )
+ patch_pos_embed = patch_pos_embed.to(torch.float)
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed,
+ scale_factor=(num_h_patches / sqrt_num_positions, num_w_patches / sqrt_num_positions),
+ mode="bicubic",
+ align_corners=False,
+ )
+ if fp32_upcasting:
+ patch_pos_embed = patch_pos_embed.to(torch.bfloat16)
+ if int(num_h_patches) != patch_pos_embed.shape[-2] or int(num_w_patches) != patch_pos_embed.shape[-1]:
+ raise ValueError(
+ f"Number of patches for images ({int(num_h_patches), int(num_w_patches)}) don't match the "
+ f"shape of position embedding ({patch_pos_embed.shape[-2], patch_pos_embed.shape[-1]})"
+ )
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, embed_dim)
+ return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+ def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = False) -> torch.Tensor:
+ batch_size, num_channels, height, width = pixel_values.shape
+ if not interpolate_pos_encoding:
+ if height != self.image_size or width != self.image_size:
+ raise ValueError(
+ f"Input image size ({height}*{width}) doesn't match model"
+ f" ({self.image_size}*{self.image_size}). You should try to set `interpolate_pos_encoding=True`"
+ )
+
+ target_dtype = self.patch_embedding.weight.dtype
+ patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid]
+
+ patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+
+ class_embeds = self.class_embedding.expand(batch_size, 1, -1)
+ embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
+
+ # add positional encoding to each token
+ if interpolate_pos_encoding:
+ embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
+ else:
+ embeddings = embeddings + self.position_embedding(self.position_ids)
+
+ return embeddings
+
+
+# Copied from transformers.models.siglip.modeling_siglip.eager_attention_forward
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+class IdeficsVisionAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config: IdeficsVisionConfig):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = self.embed_dim // self.num_heads
+ if self.head_dim * self.num_heads != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+ f" {self.num_heads})."
+ )
+ self.scale = self.head_dim**-0.5
+ self.dropout = config.attention_dropout
+ self.is_causal = False
+
+ self.k_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.v_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.q_proj = nn.Linear(self.embed_dim, self.embed_dim)
+ self.out_proj = nn.Linear(self.embed_dim, self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ """Input shape: Batch x Time x Channel"""
+
+ batch_size, seq_length, embed_dim = hidden_states.shape
+
+ queries = self.q_proj(hidden_states)
+ keys = self.k_proj(hidden_states)
+ values = self.v_proj(hidden_states)
+
+ queries = queries.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+ keys = keys.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+ values = values.view(batch_size, seq_length, self.num_heads, self.head_dim).transpose(1, 2)
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ queries,
+ keys,
+ values,
+ attention_mask,
+ is_causal=self.is_causal,
+ scaling=self.scale,
+ dropout=0.0 if not self.training else self.dropout,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(batch_size, seq_length, embed_dim).contiguous()
+ attn_output = self.out_proj(attn_output)
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.clip.modeling_clip.CLIPMLP with CLIP->IdeficsVision
+class IdeficsVisionMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.activation_fn = ACT2FN[config.hidden_act]
+ self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+ self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoderLayer with AltCLIP->IdeficsVision
+class IdeficsVisionEncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: IdeficsVisionConfig):
+ super().__init__()
+ self.embed_dim = config.hidden_size
+ self.self_attn = IdeficsVisionAttention(config)
+ self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+ self.mlp = IdeficsVisionMLP(config)
+ self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.FloatTensor, torch.Tensor | None]:
+ residual = hidden_states
+
+ hidden_states = self.layer_norm1(hidden_states)
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.layer_norm2(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+ hidden_states = residual + hidden_states
+
+ return hidden_states
+
+
+# Copied from transformers.models.altclip.modeling_altclip.AltCLIPEncoder with AltCLIP->IdeficsVision
+class IdeficsVisionEncoder(nn.Module):
+ """
+ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+ [`IdeficsVisionEncoderLayer`].
+
+ Args:
+ config: IdeficsVisionConfig
+ """
+
+ def __init__(self, config: IdeficsVisionConfig):
+ super().__init__()
+ self.config = config
+ self.layers = nn.ModuleList([IdeficsVisionEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ inputs_embeds,
+ attention_mask: torch.Tensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutput:
+ r"""
+ Args:
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ 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.
+ 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)
+
+ """
+ hidden_states = inputs_embeds
+ for encoder_layer in self.layers:
+ hidden_states = encoder_layer(
+ hidden_states,
+ attention_mask,
+ **kwargs,
+ )
+
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ )
+
+
+# Adapted from transformers.models.clip.modeling_clip.CLIPVisionTransformer
+class IdeficsVisionTransformer(nn.Module):
+ def __init__(self, config: IdeficsVisionConfig):
+ super().__init__()
+ self.config = config
+ embed_dim = config.hidden_size
+
+ self.embeddings = IdeficsVisionEmbeddings(config)
+ self.pre_layrnorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+ self.encoder = IdeficsVisionEncoder(config)
+ self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+ # Adapted from transformers.models.clip.modeling_clip.CLIPVisionTransformer.forward
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ interpolate_pos_encoding: bool | None = False,
+ **kwargs,
+ ) -> tuple | BaseModelOutputWithPooling:
+ r"""
+ Returns:
+
+ """
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
+ hidden_states = self.pre_layrnorm(hidden_states)
+
+ encoder_outputs: BaseModelOutput = self.encoder(
+ inputs_embeds=hidden_states,
+ **kwargs,
+ )
+
+ last_hidden_state = encoder_outputs.last_hidden_state
+ pooled_output = last_hidden_state[:, 0, :]
+ pooled_output = self.post_layernorm(pooled_output)
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooled_output,
+ )
diff --git a/third_party/transformers/src/transformers/models/instructblip/__init__.py b/third_party/transformers/src/transformers/models/instructblip/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ed2bd053d2c8b9fdaec4211b2b74e964bb88a5e3
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/instructblip/__init__.py
@@ -0,0 +1,28 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_instructblip import *
+ from .modeling_instructblip import *
+ from .processing_instructblip import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/instructblip/configuration_instructblip.py b/third_party/transformers/src/transformers/models/instructblip/configuration_instructblip.py
new file mode 100644
index 0000000000000000000000000000000000000000..07fe02fb613b9f439bf95933af4166dfad7d95c4
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/instructblip/configuration_instructblip.py
@@ -0,0 +1,186 @@
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""InstructBLIP model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+from ...utils import auto_docstring, logging
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="Salesforce/instructblip-flan-t5-xl")
+@strict
+class InstructBlipVisionConfig(PreTrainedConfig):
+ r"""
+ Example:
+
+ ```python
+ >>> from transformers import InstructBlipVisionConfig, InstructBlipVisionModel
+
+ >>> # Initializing a InstructBlipVisionConfig with Salesforce/instructblip-flan-t5-xl style configuration
+ >>> configuration = InstructBlipVisionConfig()
+
+ >>> # Initializing a InstructBlipVisionModel (with random weights) from the Salesforce/instructblip-flan-t5-xl style configuration
+ >>> model = InstructBlipVisionModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "instructblip_vision_model"
+ base_config_key = "vision_config"
+
+ hidden_size: int = 1408
+ intermediate_size: int = 6144
+ num_hidden_layers: int = 39
+ num_attention_heads: int = 16
+ image_size: int | list[int] | tuple[int, int] = 224
+ patch_size: int | list[int] | tuple[int, int] = 14
+ hidden_act: str = "gelu"
+ layer_norm_eps: float = 1e-6
+ attention_dropout: float | int = 0.0
+ initializer_range: float = 1e-10
+ qkv_bias: bool = True
+
+
+@auto_docstring(checkpoint="Salesforce/instructblip-flan-t5-xl")
+@strict
+class InstructBlipQFormerConfig(PreTrainedConfig):
+ r"""
+ cross_attention_frequency (`int`, *optional*, defaults to 2):
+ The frequency of adding cross-attention to the Transformer layers.
+ encoder_hidden_size (`int`, *optional*, defaults to 1408):
+ The hidden size of the hidden states for cross-attention.
+
+ Examples:
+
+ ```python
+ >>> from transformers import InstructBlipQFormerConfig, InstructBlipQFormerModel
+
+ >>> # Initializing a InstructBLIP Salesforce/instructblip-flan-t5-xl style configuration
+ >>> configuration = InstructBlipQFormerConfig()
+
+ >>> # Initializing a model (with random weights) from the Salesforce/instructblip-flan-t5-xl style configuration
+ >>> model = InstructBlipQFormerModel(configuration)
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "instructblip_qformer"
+ base_config_key = "qformer_config"
+
+ vocab_size: int = 30522
+ hidden_size: int = 768
+ num_hidden_layers: int = 12
+ num_attention_heads: int = 12
+ intermediate_size: int = 3072
+ hidden_act: str = "gelu"
+ hidden_dropout_prob: float | int = 0.1
+ attention_probs_dropout_prob: float | int = 0.1
+ max_position_embeddings: int = 512
+ initializer_range: float = 0.02
+ layer_norm_eps: float = 1e-12
+ pad_token_id: int | None = 0
+ cross_attention_frequency: int = 2
+ encoder_hidden_size: int = 1408
+
+
+@auto_docstring(checkpoint="Salesforce/instructblip-flan-t5-xl")
+@strict
+class InstructBlipConfig(PreTrainedConfig):
+ r"""
+ qformer_config (`dict`, *optional*):
+ Dictionary of configuration options used to initialize [`InstructBlipQFormerConfig`].
+ num_query_tokens (`int`, *optional*, defaults to 32):
+ The number of query tokens passed through the Transformer.
+
+ Example:
+
+ ```python
+ >>> from transformers import (
+ ... InstructBlipVisionConfig,
+ ... InstructBlipQFormerConfig,
+ ... OPTConfig,
+ ... InstructBlipConfig,
+ ... InstructBlipForConditionalGeneration,
+ ... )
+
+ >>> # Initializing a InstructBlipConfig with Salesforce/instructblip-flan-t5-xl style configuration
+ >>> configuration = InstructBlipConfig()
+
+ >>> # Initializing a InstructBlipForConditionalGeneration (with random weights) from the Salesforce/instructblip-flan-t5-xl style configuration
+ >>> model = InstructBlipForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+
+ >>> # We can also initialize a InstructBlipConfig from a InstructBlipVisionConfig, InstructBlipQFormerConfig and any PreTrainedConfig
+
+ >>> # Initializing InstructBLIP vision, InstructBLIP Q-Former and language model configurations
+ >>> vision_config = InstructBlipVisionConfig()
+ >>> qformer_config = InstructBlipQFormerConfig()
+ >>> text_config = OPTConfig()
+
+ >>> config = InstructBlipConfig(vision_config=vision_config, qformer_config=qformer_config, text_config=text_config)
+ ```"""
+
+ model_type = "instructblip"
+ attribute_map = {
+ "image_token_id": "image_token_index",
+ }
+ sub_configs = {
+ "text_config": AutoConfig,
+ "qformer_config": InstructBlipQFormerConfig,
+ "vision_config": InstructBlipVisionConfig,
+ }
+
+ vision_config: dict | PreTrainedConfig | None = None
+ qformer_config: dict | PreTrainedConfig | None = None
+ text_config: dict | PreTrainedConfig | None = None
+ num_query_tokens: int = 32
+ image_token_index: int | None = None
+ initializer_factor: float = 1.0
+ initializer_range: float = 0.02
+
+ def __post_init__(self, **kwargs):
+ if self.text_config is None:
+ self.text_config = CONFIG_MAPPING["opt"]()
+ logger.info("text_config is None. Initializing the text config with default values (`OPTConfig`).")
+ elif isinstance(self.text_config, dict):
+ text_model_type = self.text_config.get("model_type", "opt")
+ self.text_config = CONFIG_MAPPING[text_model_type](**self.text_config)
+
+ if self.qformer_config is None:
+ self.qformer_config = InstructBlipQFormerConfig()
+ logger.info("qformer_config is None. Initializing the InstructBlipQFormerConfig with default values.")
+ elif isinstance(self.qformer_config, dict):
+ self.qformer_config = InstructBlipQFormerConfig(**self.qformer_config)
+
+ if self.vision_config is None:
+ self.vision_config = InstructBlipVisionConfig()
+ logger.info("`vision_config` is `None`. initializing the `InstructBlipVisionConfig` with default values.")
+ elif isinstance(self.vision_config, dict):
+ self.vision_config = InstructBlipVisionConfig(**self.vision_config)
+
+ self.qformer_config.encoder_hidden_size = self.vision_config.hidden_size
+ self.use_decoder_only_language_model = self.text_config.model_type in MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["InstructBlipConfig", "InstructBlipQFormerConfig", "InstructBlipVisionConfig"]
diff --git a/third_party/transformers/src/transformers/models/instructblip/convert_instructblip_original_to_pytorch.py b/third_party/transformers/src/transformers/models/instructblip/convert_instructblip_original_to_pytorch.py
new file mode 100644
index 0000000000000000000000000000000000000000..8722d502d99c00036965a6364cee443dc4f83f99
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/instructblip/convert_instructblip_original_to_pytorch.py
@@ -0,0 +1,304 @@
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+"""
+Convert InstructBLIP checkpoints from the original repository.
+
+URL: https://github.com/salesforce/LAVIS/tree/main/projects/instructblip
+"""
+
+import argparse
+from io import BytesIO
+
+import httpx
+import torch
+
+# pip3 install salesforce-lavis
+# I'm actually installing a slightly modified version: pip3 install git+https://github.com/nielsrogge/LAVIS.git@fix_lavis_float32 (there's also the fix_lavis branch)
+# also note: to convert Vicuna checkpoints, we had to include /home/niels/python_projects/checkpoints/FastChat/vicuna-7b in lavis/configs/models/blip2/blip2_instruct_vicuna7b.yaml
+# same for Vicuna-13b
+from lavis.models import load_model_and_preprocess
+from PIL import Image
+
+from transformers import (
+ AutoTokenizer,
+ BlipImageProcessor,
+ InstructBlipConfig,
+ InstructBlipForConditionalGeneration,
+ InstructBlipProcessor,
+ InstructBlipQFormerConfig,
+ InstructBlipVisionConfig,
+ LlamaConfig,
+ LlamaTokenizerFast,
+ T5Config,
+ T5TokenizerFast,
+)
+from transformers.utils.constants import OPENAI_CLIP_MEAN, OPENAI_CLIP_STD
+
+
+def load_demo_image():
+ url = "https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"
+ with httpx.stream("GET", url) as response:
+ image = Image.open(BytesIO(response.read())).convert("RGB")
+
+ return image
+
+
+# here we list all keys to be renamed (original name on the left, our name on the right)
+def create_rename_keys(config):
+ rename_keys = []
+ # fmt: off
+
+ # vision encoder
+ rename_keys.append(("visual_encoder.cls_token", "vision_model.embeddings.class_embedding"))
+ rename_keys.append(("visual_encoder.pos_embed", "vision_model.embeddings.position_embedding"))
+ rename_keys.append(("visual_encoder.patch_embed.proj.weight", "vision_model.embeddings.patch_embedding.weight"))
+ rename_keys.append(("visual_encoder.patch_embed.proj.bias", "vision_model.embeddings.patch_embedding.bias"))
+ rename_keys.append(("ln_vision.weight", "vision_model.post_layernorm.weight"))
+ rename_keys.append(("ln_vision.bias", "vision_model.post_layernorm.bias"))
+
+ for i in range(config.vision_config.num_hidden_layers):
+ rename_keys.append((f"visual_encoder.blocks.{i}.norm1.weight", f"vision_model.encoder.layers.{i}.layer_norm1.weight"))
+ rename_keys.append((f"visual_encoder.blocks.{i}.norm1.bias", f"vision_model.encoder.layers.{i}.layer_norm1.bias"))
+ rename_keys.append((f"visual_encoder.blocks.{i}.norm2.weight", f"vision_model.encoder.layers.{i}.layer_norm2.weight"))
+ rename_keys.append((f"visual_encoder.blocks.{i}.norm2.bias", f"vision_model.encoder.layers.{i}.layer_norm2.bias"))
+ rename_keys.append((f"visual_encoder.blocks.{i}.attn.qkv.weight", f"vision_model.encoder.layers.{i}.self_attn.qkv.weight"))
+ rename_keys.append((f"visual_encoder.blocks.{i}.attn.proj.weight", f"vision_model.encoder.layers.{i}.self_attn.projection.weight",))
+ rename_keys.append((f"visual_encoder.blocks.{i}.attn.proj.bias", f"vision_model.encoder.layers.{i}.self_attn.projection.bias"))
+ rename_keys.append((f"visual_encoder.blocks.{i}.mlp.fc1.weight", f"vision_model.encoder.layers.{i}.mlp.fc1.weight"))
+ rename_keys.append((f"visual_encoder.blocks.{i}.mlp.fc1.bias", f"vision_model.encoder.layers.{i}.mlp.fc1.bias"))
+ rename_keys.append((f"visual_encoder.blocks.{i}.mlp.fc2.weight", f"vision_model.encoder.layers.{i}.mlp.fc2.weight"))
+ rename_keys.append((f"visual_encoder.blocks.{i}.mlp.fc2.bias", f"vision_model.encoder.layers.{i}.mlp.fc2.bias"))
+
+ # QFormer
+ rename_keys.append(("Qformer.bert.embeddings.LayerNorm.weight", "qformer.embeddings.layernorm.weight"))
+ rename_keys.append(("Qformer.bert.embeddings.LayerNorm.bias", "qformer.embeddings.layernorm.bias"))
+
+ # fmt: on
+ return rename_keys
+
+
+def rename_key(dct, old, new):
+ val = dct.pop(old)
+ dct[new] = val
+
+
+def read_in_q_v_bias(state_dict, config):
+ for i in range(config.vision_config.num_hidden_layers):
+ # read in original q and v biases
+ q_bias = state_dict.pop(f"visual_encoder.blocks.{i}.attn.q_bias")
+ v_bias = state_dict.pop(f"visual_encoder.blocks.{i}.attn.v_bias")
+
+ # next, set bias in the state dict
+ qkv_bias = torch.cat((q_bias, torch.zeros_like(v_bias, requires_grad=False), v_bias))
+ state_dict[f"vision_model.encoder.layers.{i}.self_attn.qkv.bias"] = qkv_bias
+
+
+def get_blip2_config(model_name):
+ image_size = 364 if "coco" in model_name else 224
+ vision_config = InstructBlipVisionConfig(image_size=image_size).to_dict()
+
+ # make sure the models have proper bos_token_id and eos_token_id set (important for generation)
+ # seems like flan-T5 models don't have bos_token_id properly set?
+ if "t5-xl" in model_name:
+ text_config = T5Config.from_pretrained("google/flan-t5-xl", dense_act_fn="gelu", bos_token_id=1).to_dict()
+ elif "t5-xxl" in model_name:
+ text_config = T5Config.from_pretrained("google/flan-t5-xxl", dense_act_fn="gelu", bos_token_id=1).to_dict()
+ elif "vicuna-7b" in model_name:
+ text_config = LlamaConfig.from_pretrained("decapoda-research/llama-7b-hf", vocab_size=32001).to_dict()
+ elif "vicuna-13b" in model_name:
+ text_config = LlamaConfig.from_pretrained("decapoda-research/llama-13b-hf", vocab_size=32001).to_dict()
+ else:
+ raise ValueError("Model name not supported")
+
+ # the authors add one special "[DEC]" token to the vocab of Q-Former, hence vocab size = 30522 + 1
+ qformer_config = InstructBlipQFormerConfig(vocab_size=30523).to_dict()
+ config = InstructBlipConfig(vision_config=vision_config, text_config=text_config, qformer_config=qformer_config)
+
+ return config, image_size
+
+
+@torch.no_grad()
+def convert_blip2_checkpoint(model_name, pytorch_dump_folder_path=None, push_to_hub=False):
+ """
+ Copy/paste/tweak model's weights to Transformers design.
+ """
+ qformer_tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased", truncation_side="left")
+ qformer_tokenizer.add_special_tokens({"bos_token": "[DEC]"})
+
+ if "t5" in model_name:
+ tokenizer = T5TokenizerFast.from_pretrained("google/flan-t5-xl", truncation_side="left")
+ elif "vicuna" in model_name:
+ # the following was used in the original implementation:
+ # tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b", use_fast=False, truncation_side="left")
+ # tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+ # tokenizer.add_special_tokens({"bos_token": ""})
+ # tokenizer.add_special_tokens({"eos_token": ""})
+ # tokenizer.add_special_tokens({"unk_token": ""})
+ tokenizer = LlamaTokenizerFast.from_pretrained(
+ "huggyllama/llama-7b", truncation_side="left", bos_token="", unk_token=""
+ )
+ tokenizer.add_special_tokens({"pad_token": "[PAD]"})
+
+ config, image_size = get_blip2_config(model_name)
+ hf_model = InstructBlipForConditionalGeneration(config).eval()
+
+ model_name_to_original = {
+ "instructblip-vicuna-7b": ("blip2_vicuna_instruct", "vicuna7b"),
+ "instructblip-vicuna-13b": ("blip2_vicuna_instruct", "vicuna13b"),
+ "instructblip-flan-t5-xl": ("blip2_t5_instruct", "flant5xl"),
+ "instructblip-flan-t5-xxl": ("blip2_t5_instruct", "flant5xxl"),
+ }
+
+ name, type = model_name_to_original[model_name]
+
+ # load original model
+ print("Loading original model...")
+ hf_model_device = "cuda:1" if torch.cuda.is_available() else "cpu"
+ lavis_device = "cuda:2" if torch.cuda.is_available() else "cpu"
+ original_model, vis_processors, _ = load_model_and_preprocess(
+ name=name, model_type=type, is_eval=True, device=lavis_device
+ )
+ original_model.eval()
+ print("Done!")
+
+ # update state dict keys
+ state_dict = original_model.state_dict()
+ rename_keys = create_rename_keys(config)
+ for src, dest in rename_keys:
+ rename_key(state_dict, src, dest)
+
+ # some keys can be renamed efficiently
+ for key, val in state_dict.copy().items():
+ val = state_dict.pop(key)
+ if key.startswith("Qformer.bert"):
+ key = key.replace("Qformer.bert", "qformer")
+ if "attention.self" in key:
+ key = key.replace("self", "attention")
+ if "llm_proj" in key:
+ key = key.replace("llm_proj", "language_projection")
+ if "t5_proj" in key:
+ key = key.replace("t5_proj", "language_projection")
+ if key.startswith("llm_model"):
+ key = key.replace("llm_model", "language_model")
+ if key.startswith("t5"):
+ key = key.replace("t5", "language")
+ state_dict[key] = val
+
+ # read in qv biases
+ read_in_q_v_bias(state_dict, config)
+
+ # note: weights get loaded in torch.float32 by default
+ hf_model.load_state_dict(state_dict, strict=True)
+
+ image = load_demo_image()
+ prompt = "What is unusual about this image?"
+
+ # create processor
+ image_processor = BlipImageProcessor(
+ size={"height": image_size, "width": image_size}, image_mean=OPENAI_CLIP_MEAN, image_std=OPENAI_CLIP_STD
+ )
+ processor = InstructBlipProcessor(
+ image_processor=image_processor,
+ tokenizer=tokenizer,
+ qformer_tokenizer=qformer_tokenizer,
+ )
+ inputs = processor(images=image, text=prompt, return_tensors="pt").to(hf_model_device)
+
+ # make sure processor creates exact same pixel values
+ original_pixel_values = vis_processors["eval"](image).unsqueeze(0).to(lavis_device)
+ pixel_values = inputs.pixel_values
+ assert torch.allclose(original_pixel_values.to(pixel_values.device), pixel_values)
+
+ original_model.to(lavis_device)
+ hf_model.to(hf_model_device)
+ with torch.no_grad():
+ if "vicuna" in model_name:
+ original_logits = original_model({"image": original_pixel_values, "text_input": [prompt]}).logits
+ logits = hf_model(**inputs).logits
+ else:
+ original_logits = original_model(
+ {"image": original_pixel_values, "text_input": [prompt], "text_output": ["\n"]}
+ ).logits
+ label_input_ids = tokenizer("\n", return_tensors="pt").input_ids.to(hf_model_device)
+ labels = label_input_ids.masked_fill(label_input_ids == tokenizer.pad_token_id, -100)
+ logits = hf_model(**inputs, labels=labels).logits
+
+ print("First values of original logits:", original_logits[0, :3, :3])
+ print("First values of HF logits:", logits[0, :3, :3])
+
+ # assert values
+ assert original_logits.shape == logits.shape
+ atol = 1e-4 if "vicuna" in model_name else 1e-5
+ assert torch.allclose(original_logits.to(logits.device), logits, atol=atol)
+ print("Looks ok!")
+
+ print("Generating with original model...")
+ original_outputs = original_model.generate({"image": original_pixel_values, "prompt": prompt}, num_beams=5)
+
+ # important: we need to cast the weights of the HF model to the appropriate type
+ print("Generating with HF model...")
+ outputs = hf_model.generate(
+ **inputs,
+ do_sample=False,
+ num_beams=5,
+ max_length=256,
+ min_length=1,
+ top_p=0.9,
+ repetition_penalty=1.5,
+ length_penalty=1.0,
+ temperature=1,
+ )
+ if "vicuna" in model_name:
+ # convert output id 0 to 2 (eos_token_id)
+ # TODO add this in the generate method?
+ outputs[outputs == 0] = 2
+ print("Original generation:", original_outputs)
+ output_text = processor.batch_decode(outputs, skip_special_tokens=True)
+ output_text = [text.strip() for text in output_text]
+ print("HF generation:", output_text)
+
+ if pytorch_dump_folder_path is not None:
+ processor.save_pretrained(pytorch_dump_folder_path)
+ hf_model.save_pretrained(pytorch_dump_folder_path)
+
+ if push_to_hub:
+ processor.push_to_hub(f"Salesforce/{model_name}")
+ hf_model.push_to_hub(f"Salesforce/{model_name}")
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ choices = [
+ "instructblip-vicuna-7b",
+ "instructblip-vicuna-13b",
+ "instructblip-flan-t5-xl",
+ "instructblip-flan-t5-xxl",
+ ]
+ parser.add_argument(
+ "--model_name",
+ default="instructblip-flan-t5-xl",
+ choices=choices,
+ type=str,
+ help="Path to hf config.json of model to convert",
+ )
+ parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
+ parser.add_argument(
+ "--push_to_hub",
+ action="store_true",
+ help="Whether to push the model and processor to the hub after converting",
+ )
+
+ args = parser.parse_args()
+
+ convert_blip2_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)
diff --git a/third_party/transformers/src/transformers/models/instructblip/modeling_instructblip.py b/third_party/transformers/src/transformers/models/instructblip/modeling_instructblip.py
new file mode 100644
index 0000000000000000000000000000000000000000..29f32f17d6c4667cc8b09ddb99cf9482821140a1
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/instructblip/modeling_instructblip.py
@@ -0,0 +1,1475 @@
+# Copyright 2023 The Salesforce Authors and The HuggingFace Team. All rights reserved.
+#
+# 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 InstructBLIP model."""
+
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...generation import GenerationMixin
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import (
+ BaseModelOutput,
+ BaseModelOutputWithPastAndCrossAttentions,
+ BaseModelOutputWithPooling,
+ BaseModelOutputWithPoolingAndCrossAttentions,
+ CausalLMOutputWithPast,
+ Seq2SeqLMOutput,
+)
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...pytorch_utils import apply_chunking_to_forward
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_int
+from ...utils.generic import merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from ..auto import AutoModel, AutoModelForCausalLM, AutoModelForSeq2SeqLM
+from .configuration_instructblip import InstructBlipConfig, InstructBlipQFormerConfig, InstructBlipVisionConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@dataclass
+@auto_docstring
+class BaseModelOutputWithVisionQformerOutputs(BaseModelOutputWithPooling):
+ r"""
+ vision_outputs (`BaseModelOutputWithPooling`):
+ Outputs of the vision encoder.
+ qformer_outputs (`BaseModelOutputWithPoolingAndCrossAttentions`):
+ Outputs of the Q-Former (Querying Transformer).
+ """
+
+ vision_outputs: BaseModelOutputWithPooling | None = None
+ qformer_outputs: BaseModelOutputWithPoolingAndCrossAttentions | None = None
+
+
+@dataclass
+@auto_docstring(
+ custom_intro="""
+ Class defining the outputs of [`InstructBlipForConditionalGeneration`].
+ """
+)
+# Copied from transformers.models.blip_2.modeling_blip_2.Blip2ForConditionalGenerationModelOutput with Blip2->InstructBlip
+class InstructBlipForConditionalGenerationModelOutput(ModelOutput):
+ r"""
+ loss (`torch.FloatTensor`, *optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
+ Language modeling loss from the language model.
+ logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+ Prediction scores of the language modeling head of the language model.
+ vision_outputs (`BaseModelOutputWithPooling`):
+ Outputs of the vision encoder.
+ qformer_outputs (`BaseModelOutputWithPoolingAndCrossAttentions`):
+ Outputs of the Q-Former (Querying Transformer).
+ language_model_outputs (`CausalLMOutputWithPast` or `Seq2SeqLMOutput`):
+ Outputs of the language model.
+ """
+
+ loss: tuple[torch.FloatTensor] | None = None
+ logits: tuple[torch.FloatTensor] | None = None
+ vision_outputs: BaseModelOutputWithPooling | None = None
+ qformer_outputs: BaseModelOutputWithPoolingAndCrossAttentions | None = None
+ language_model_outputs: CausalLMOutputWithPast | Seq2SeqLMOutput | None = None
+
+ def to_tuple(self) -> tuple[Any]:
+ return tuple(
+ self[k]
+ if k not in ["vision_outputs", "qformer_outputs", "language_model_outputs"]
+ else getattr(self, k).to_tuple()
+ for k in self.keys()
+ )
+
+
+# Copied from transformers.models.blip.modeling_blip.BlipVisionEmbeddings with Blip->InstructBlip
+class InstructBlipVisionEmbeddings(nn.Module):
+ def __init__(self, config: InstructBlipVisionConfig):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.image_size = config.image_size
+ self.patch_size = config.patch_size
+
+ self.class_embedding = nn.Parameter(torch.randn(1, 1, self.embed_dim))
+
+ self.patch_embedding = nn.Conv2d(
+ in_channels=3, out_channels=self.embed_dim, kernel_size=self.patch_size, stride=self.patch_size
+ )
+
+ self.num_patches = (self.image_size // self.patch_size) ** 2
+ self.num_positions = self.num_patches + 1
+
+ self.position_embedding = nn.Parameter(torch.randn(1, self.num_positions, self.embed_dim))
+
+ def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
+ """
+ This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher resolution
+ images. This method is also adapted to support torch.jit tracing.
+
+ Adapted from:
+ - https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174-L194, and
+ - https://github.com/facebookresearch/dinov2/blob/e1277af2ba9496fbadf7aec6eba56e8d882d1e35/dinov2/models/vision_transformer.py#L179-L211
+ """
+
+ num_patches = embeddings.shape[1] - 1
+ num_positions = self.position_embedding.shape[1] - 1
+
+ # always interpolate when tracing to ensure the exported model works for dynamic input shapes
+ if not torch.jit.is_tracing() and num_patches == num_positions and height == width:
+ return self.position_embedding
+
+ class_pos_embed = self.position_embedding[:, :1]
+ patch_pos_embed = self.position_embedding[:, 1:]
+
+ dim = embeddings.shape[-1]
+
+ new_height = height // self.patch_size
+ new_width = width // self.patch_size
+
+ sqrt_num_positions = torch_int(num_positions**0.5)
+ patch_pos_embed = patch_pos_embed.reshape(1, sqrt_num_positions, sqrt_num_positions, dim)
+ patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
+
+ patch_pos_embed = nn.functional.interpolate(
+ patch_pos_embed,
+ size=(new_height, new_width),
+ mode="bicubic",
+ align_corners=False,
+ )
+
+ patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+
+ return torch.cat((class_pos_embed, patch_pos_embed), dim=1)
+
+ def forward(self, pixel_values: torch.FloatTensor, interpolate_pos_encoding: bool = False) -> torch.Tensor:
+ batch_size, _, height, width = pixel_values.shape
+ target_dtype = self.patch_embedding.weight.dtype
+ patch_embeds = self.patch_embedding(pixel_values.to(dtype=target_dtype)) # shape = [*, width, grid, grid]
+ patch_embeds = patch_embeds.flatten(2).transpose(1, 2)
+ class_embeds = self.class_embedding.expand(batch_size, 1, -1).to(target_dtype)
+ embeddings = torch.cat([class_embeds, patch_embeds], dim=1)
+ if interpolate_pos_encoding:
+ position_embedding = self.interpolate_pos_encoding(embeddings, height, width)
+ else:
+ position_embedding = self.position_embedding
+ embeddings = embeddings + position_embedding[:, : embeddings.size(1), :].to(target_dtype)
+ return embeddings
+
+
+# Adapted from transformers.models.siglip.modeling_siglip.eager_attention_forward -> InstructBLIP doesn't cast attn weights to fp32
+def eager_attention_forward(
+ module: nn.Module,
+ query: torch.Tensor,
+ key: torch.Tensor,
+ value: torch.Tensor,
+ attention_mask: torch.Tensor | None,
+ scaling: float,
+ dropout: float = 0.0,
+ **kwargs,
+):
+ attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
+ if attention_mask is not None:
+ attn_weights = attn_weights + attention_mask
+
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+ attn_output = torch.matmul(attn_weights, value)
+ attn_output = attn_output.transpose(1, 2).contiguous()
+
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.blip_2.modeling_blip_2.Blip2Attention with Blip2->InstructBlip
+class InstructBlipAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.embed_dim = config.hidden_size
+ self.num_heads = config.num_attention_heads
+ self.head_dim = self.embed_dim // self.num_heads
+ if self.head_dim * self.num_heads != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
+ f" {self.num_heads})."
+ )
+ self.scale = self.head_dim**-0.5
+ self.is_causal = False
+ self.attention_dropout = config.attention_dropout
+
+ # small tweak here compared to CLIP, no bias here
+ self.qkv = nn.Linear(self.embed_dim, 3 * self.embed_dim, bias=False)
+
+ if config.qkv_bias:
+ q_bias = nn.Parameter(torch.zeros(self.embed_dim))
+ v_bias = nn.Parameter(torch.zeros(self.embed_dim))
+ else:
+ q_bias = None
+ v_bias = None
+
+ if q_bias is not None:
+ qkv_bias = torch.cat((q_bias, torch.zeros_like(v_bias, requires_grad=False), v_bias))
+ self.qkv.bias = nn.Parameter(qkv_bias)
+
+ self.projection = nn.Linear(self.embed_dim, self.embed_dim)
+
+ 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,
+ **kwargs,
+ ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+ """Input shape: Batch x Time x Channel"""
+
+ bsz, tgt_len, embed_dim = hidden_states.size()
+
+ mixed_qkv = self.qkv(hidden_states)
+
+ mixed_qkv = mixed_qkv.reshape(bsz, tgt_len, 3, self.num_heads, embed_dim // self.num_heads).permute(
+ 2, 0, 3, 1, 4
+ )
+ query_states, key_states, value_states = mixed_qkv[0], mixed_qkv[1], mixed_qkv[2]
+
+ attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+ self.config._attn_implementation, eager_attention_forward
+ )
+
+ attn_output, attn_weights = attention_interface(
+ self,
+ query_states,
+ key_states,
+ value_states,
+ attention_mask=None,
+ dropout=0.0 if not self.training else self.attention_dropout,
+ scaling=self.scale,
+ **kwargs,
+ )
+
+ attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()
+ attn_output = self.projection(attn_output)
+
+ return attn_output, attn_weights
+
+
+# Copied from transformers.models.blip.modeling_blip.BlipMLP
+class InstructBlipMLP(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.activation_fn = ACT2FN[config.hidden_act]
+ self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size)
+ self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.fc1(hidden_states)
+ hidden_states = self.activation_fn(hidden_states)
+ hidden_states = self.fc2(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.blip.modeling_blip.BlipEncoderLayer with Blip->InstructBlip
+class InstructBlipEncoderLayer(GradientCheckpointingLayer):
+ def __init__(self, config: InstructBlipConfig):
+ super().__init__()
+ self.embed_dim = config.hidden_size
+ self.self_attn = InstructBlipAttention(config)
+ self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+ self.mlp = InstructBlipMLP(config)
+ self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps)
+
+ @auto_docstring
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.FloatTensor:
+ residual = hidden_states
+
+ hidden_states = self.layer_norm1(hidden_states)
+ hidden_states, _ = self.self_attn(
+ hidden_states=hidden_states,
+ **kwargs,
+ )
+ hidden_states = hidden_states + residual
+ residual = hidden_states
+ hidden_states = self.layer_norm2(hidden_states)
+ hidden_states = self.mlp(hidden_states)
+
+ hidden_states = hidden_states + residual
+
+ return hidden_states
+
+
+@auto_docstring
+class InstructBlipPreTrainedModel(PreTrainedModel):
+ config: InstructBlipConfig
+ base_model_prefix = "blip"
+ input_modalities = ("image", "text")
+ supports_gradient_checkpointing = True
+ _supports_attention_backend = True
+ _supports_flash_attn = True
+ _supports_sdpa = True
+ _supports_flex_attn = True
+
+ _can_compile_fullgraph = True
+
+ _no_split_modules = [
+ "InstructBlipQFormerEmbeddings",
+ "InstructBlipAttention",
+ "InstructBlipQFormerMultiHeadAttention",
+ "InstructBlipQFormerSelfOutput",
+ ]
+
+ @torch.no_grad()
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ super()._init_weights(module)
+ factor = self.config.initializer_range
+ if isinstance(module, InstructBlipVisionEmbeddings):
+ init.trunc_normal_(module.position_embedding, mean=0.0, std=factor)
+ init.trunc_normal_(module.class_embedding, mean=0.0, std=factor)
+ elif isinstance(module, (InstructBlipForConditionalGeneration, InstructBlipModel)):
+ init.zeros_(module.query_tokens)
+ elif isinstance(module, InstructBlipQFormerEmbeddings):
+ init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)))
+
+
+# Copied from transformers.models.blip.modeling_blip.BlipEncoder with Blip->InstructBlip
+class InstructBlipEncoder(nn.Module):
+ """
+ Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a
+ [`InstructBlipEncoderLayer`].
+
+ Args:
+ config (`InstructBlipConfig`):
+ The corresponding vision configuration for the `InstructBlipEncoder`.
+ """
+
+ def __init__(self, config: InstructBlipConfig):
+ super().__init__()
+ self.config = config
+ self.layers = nn.ModuleList([InstructBlipEncoderLayer(config) for _ in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+
+ @auto_docstring
+ def forward(
+ self,
+ inputs_embeds,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutput:
+ hidden_states = inputs_embeds
+ for encoder_layer in self.layers:
+ hidden_states = encoder_layer(
+ hidden_states,
+ **kwargs,
+ )
+
+ return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+class InstructBlipVisionModel(InstructBlipPreTrainedModel):
+ main_input_name = "pixel_values"
+ input_modalities = ("image",)
+ config: InstructBlipVisionConfig
+ _can_record_outputs = {
+ "hidden_states": InstructBlipEncoderLayer,
+ "attentions": InstructBlipAttention,
+ }
+
+ def __init__(self, config: InstructBlipVisionConfig):
+ super().__init__(config)
+ self.config = config
+ embed_dim = config.hidden_size
+
+ self.embeddings = InstructBlipVisionEmbeddings(config)
+ self.encoder = InstructBlipEncoder(config)
+ self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)
+
+ self.post_init()
+
+ @merge_with_config_defaults
+ @capture_outputs(tie_last_hidden_states=False)
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithPooling:
+ if pixel_values is None:
+ raise ValueError("You have to specify pixel_values")
+
+ hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
+
+ encoder_outputs: BaseModelOutput = self.encoder(
+ inputs_embeds=hidden_states,
+ **kwargs,
+ )
+
+ last_hidden_state = encoder_outputs.last_hidden_state
+ last_hidden_state = self.post_layernorm(last_hidden_state)
+
+ pooled_output = last_hidden_state[:, 0, :]
+ pooled_output = self.post_layernorm(pooled_output)
+
+ return BaseModelOutputWithPooling(
+ last_hidden_state=last_hidden_state,
+ pooler_output=pooled_output,
+ )
+
+ def get_input_embeddings(self):
+ return self.embeddings
+
+
+class InstructBlipQFormerMultiHeadAttention(nn.Module):
+ def __init__(self, config, is_cross_attention=False):
+ super().__init__()
+ self.config = config
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+ raise ValueError(
+ "The hidden size (%d) is not a multiple of the number of attention heads (%d)"
+ % (config.hidden_size, config.num_attention_heads)
+ )
+
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
+ if is_cross_attention:
+ self.key = nn.Linear(config.encoder_hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.encoder_hidden_size, self.all_head_size)
+ else:
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+ self.save_attention = False
+
+ def save_attn_gradients(self, attn_gradients):
+ self.attn_gradients = attn_gradients
+
+ def get_attn_gradients(self):
+ return self.attn_gradients
+
+ def save_attention_map(self, attention_map):
+ self.attention_map = attention_map
+
+ def get_attention_map(self):
+ return self.attention_map
+
+ def transpose_for_scores(self, x):
+ new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+ x = x.view(*new_x_shape)
+ return x.permute(0, 2, 1, 3)
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ # If this is instantiated as a cross-attention module, the keys
+ # and values come from an encoder; the attention mask needs to be
+ # such that the encoder's padding tokens are not attended to.
+ is_cross_attention = encoder_hidden_states is not None
+
+ if is_cross_attention:
+ key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
+ value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
+ attention_mask = encoder_attention_mask
+ else:
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
+
+ mixed_query_layer = self.query(hidden_states)
+
+ query_layer = self.transpose_for_scores(mixed_query_layer)
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+ attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+ attention_scores_dtype = attention_scores.dtype
+
+ if attention_mask is not None:
+ # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
+ attention_scores = attention_scores + attention_mask
+
+ # Normalize the attention scores to probabilities.
+ attention_probs = nn.Softmax(dim=-1)(attention_scores).to(attention_scores_dtype)
+
+ if is_cross_attention and self.save_attention:
+ self.save_attention_map(attention_probs)
+ attention_probs.register_hook(self.save_attn_gradients)
+
+ # This is actually dropping out entire tokens to attend to, which might
+ # seem a bit unusual, but is taken from the original Transformer paper.
+ attention_probs_dropped = self.dropout(attention_probs)
+
+ context_layer = torch.matmul(attention_probs_dropped, value_layer)
+
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+ context_layer = context_layer.view(*new_context_layer_shape)
+
+ return context_layer, attention_probs
+
+
+# Copied from transformers.models.bert.modeling_bert.BertSelfOutput with Bert->InstructBlipQFormer
+class InstructBlipQFormerSelfOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+# Copied from transformers.models.blip_2.modeling_blip_2.Blip2QFormerAttention with Blip2->InstructBlip
+class InstructBlipQFormerAttention(nn.Module):
+ def __init__(self, config, is_cross_attention=False):
+ super().__init__()
+ self.attention = InstructBlipQFormerMultiHeadAttention(config, is_cross_attention)
+ self.output = InstructBlipQFormerSelfOutput(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: torch.FloatTensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> torch.Tensor:
+ attn_output, _ = self.attention(
+ hidden_states=hidden_states,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ **kwargs,
+ )
+ attention_output = self.output(attn_output, hidden_states)
+ return attention_output
+
+
+# Copied from transformers.models.bert.modeling_bert.BertIntermediate with Bert->InstructBlipQFormer
+class InstructBlipQFormerIntermediate(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+# Copied from transformers.models.bert.modeling_bert.BertOutput with Bert->InstructBlipQFormer
+class InstructBlipQFormerOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+class InstructBlipQFormerLayer(GradientCheckpointingLayer):
+ def __init__(self, config, layer_idx):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+ self.attention = InstructBlipQFormerAttention(config)
+
+ self.layer_idx = layer_idx
+
+ if layer_idx % config.cross_attention_frequency == 0:
+ self.crossattention = InstructBlipQFormerAttention(config, is_cross_attention=True)
+ self.has_cross_attention = True
+ else:
+ self.has_cross_attention = False
+
+ self.intermediate = InstructBlipQFormerIntermediate(config)
+ self.output = InstructBlipQFormerOutput(config)
+
+ self.intermediate_query = InstructBlipQFormerIntermediate(config)
+ self.output_query = InstructBlipQFormerOutput(config)
+
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ query_length=0,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ attention_output = self.attention(
+ hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+
+ if query_length > 0:
+ query_attention_output = attention_output[:, :query_length, :]
+
+ if self.has_cross_attention:
+ if encoder_hidden_states is None:
+ raise ValueError("encoder_hidden_states must be given for cross-attention layers")
+ query_attention_output = self.crossattention(
+ query_attention_output,
+ attention_mask=attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_attention_mask,
+ **kwargs,
+ )
+
+ layer_output = apply_chunking_to_forward(
+ self.feed_forward_chunk_query,
+ self.chunk_size_feed_forward,
+ self.seq_len_dim,
+ query_attention_output,
+ )
+
+ if attention_output.shape[1] > query_length:
+ layer_output_text = apply_chunking_to_forward(
+ self.feed_forward_chunk,
+ self.chunk_size_feed_forward,
+ self.seq_len_dim,
+ attention_output[:, query_length:, :],
+ ).to(layer_output.device)
+ layer_output = torch.cat([layer_output, layer_output_text], dim=1)
+ else:
+ layer_output = apply_chunking_to_forward(
+ self.feed_forward_chunk,
+ self.chunk_size_feed_forward,
+ self.seq_len_dim,
+ attention_output,
+ )
+ return layer_output
+
+ def feed_forward_chunk(self, attention_output):
+ intermediate_output = self.intermediate(attention_output)
+ layer_output = self.output(intermediate_output, attention_output)
+ return layer_output
+
+ def feed_forward_chunk_query(self, attention_output):
+ intermediate_output = self.intermediate_query(attention_output)
+ layer_output = self.output_query(intermediate_output, attention_output)
+ return layer_output
+
+
+# Copied from transformers.models.blip_2.modeling_blip_2.Blip2QFormerEncoder with Blip2->InstructBlip
+class InstructBlipQFormerEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList(
+ [InstructBlipQFormerLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
+ )
+ self.gradient_checkpointing = False
+
+ @can_return_tuple
+ def forward(
+ self,
+ hidden_states,
+ attention_mask=None,
+ encoder_hidden_states=None,
+ encoder_attention_mask=None,
+ query_length=0,
+ **kwargs: Unpack[TransformersKwargs],
+ ):
+ for i in range(self.config.num_hidden_layers):
+ layer_module = self.layer[i]
+
+ hidden_states = layer_module(
+ hidden_states,
+ attention_mask,
+ encoder_hidden_states, # as a positional argument for gradient checkpointing
+ encoder_attention_mask=encoder_attention_mask,
+ query_length=query_length,
+ **kwargs,
+ )
+
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ )
+
+
+class InstructBlipQFormerEmbeddings(nn.Module):
+ """Construct the embeddings from word and position embeddings."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+ self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+
+ self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+ self.register_buffer(
+ "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+ )
+
+ self.config = config
+
+ def forward(
+ self,
+ input_ids=None,
+ position_ids=None,
+ query_embeds=None,
+ past_key_values_length=0,
+ ):
+ if input_ids is not None:
+ seq_length = input_ids.size()[1]
+ else:
+ seq_length = 0
+
+ if position_ids is None:
+ position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length].clone()
+
+ if input_ids is not None:
+ embeddings = self.word_embeddings(input_ids)
+
+ position_embeddings = self.position_embeddings(position_ids.to(embeddings.device))
+ embeddings = embeddings + position_embeddings
+
+ if query_embeds is not None:
+ embeddings = torch.cat((query_embeds, embeddings), dim=1)
+ else:
+ embeddings = query_embeds
+
+ embeddings = embeddings.to(self.layernorm.weight.dtype)
+ embeddings = self.layernorm(embeddings)
+ embeddings = self.dropout(embeddings)
+ return embeddings
+
+
+class InstructBlipQFormerModel(InstructBlipPreTrainedModel):
+ """
+ Querying Transformer (Q-Former), used in InstructBLIP. Slightly modified from BLIP-2 as it also takes the
+ instruction as input.
+ """
+
+ _supports_attention_backend = False # adds position on attn weights before last matmul
+ _supports_flash_attn = False
+ _supports_sdpa = False
+ _supports_flex_attn = False
+
+ _can_record_outputs = {
+ "hidden_states": InstructBlipQFormerLayer,
+ "attentions": [
+ OutputRecorder(InstructBlipQFormerMultiHeadAttention, index=1, layer_name=".attention"),
+ ],
+ "cross_attentions": [
+ OutputRecorder(InstructBlipQFormerMultiHeadAttention, index=1, layer_name=".crossattention"),
+ ],
+ }
+
+ def __init__(self, config: InstructBlipQFormerConfig):
+ super().__init__(config)
+ self.config = config
+
+ self.embeddings = InstructBlipQFormerEmbeddings(config)
+
+ self.encoder = InstructBlipQFormerEncoder(config)
+
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embeddings.word_embeddings
+
+ def set_input_embeddings(self, value):
+ self.embeddings.word_embeddings = value
+
+ def get_extended_attention_mask(
+ self,
+ attention_mask: torch.Tensor,
+ input_shape: tuple[int],
+ device: torch.device,
+ has_query: bool = False,
+ ) -> torch.Tensor:
+ """
+ Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
+
+ Arguments:
+ attention_mask (`torch.Tensor`):
+ Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
+ input_shape (`tuple[int]`):
+ The shape of the input to the model.
+ device: (`torch.device`):
+ The device of the input to the model.
+
+ Returns:
+ `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.
+ """
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+ # ourselves in which case we just need to make it broadcastable to all heads.
+ if attention_mask.dim() == 3:
+ extended_attention_mask = attention_mask[:, None, :, :]
+ elif attention_mask.dim() == 2:
+ # Provided a padding mask of dimensions [batch_size, seq_length]
+ # - the model is an encoder, so make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
+ extended_attention_mask = attention_mask[:, None, None, :]
+ else:
+ raise ValueError(
+ f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})",
+ )
+
+ # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+ # masked positions, this operation will create a tensor which is 0.0 for
+ # positions we want to attend and -10000.0 for masked positions.
+ # Since we are adding it to the raw scores before the softmax, this is
+ # effectively the same as removing these entirely.
+ extended_attention_mask = extended_attention_mask.to(dtype=self.dtype) # fp16 compatibility
+ extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+ return extended_attention_mask
+
+ @merge_with_config_defaults
+ @capture_outputs
+ @auto_docstring
+ def forward(
+ self,
+ input_ids: torch.LongTensor,
+ attention_mask: torch.FloatTensor | None = None,
+ position_ids: torch.LongTensor | None = None,
+ query_embeds: torch.Tensor | None = None,
+ encoder_hidden_states: torch.FloatTensor | None = None,
+ encoder_attention_mask: torch.FloatTensor | None = None,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple[torch.FloatTensor] | BaseModelOutputWithPoolingAndCrossAttentions:
+ r"""
+ query_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ Hidden states to be used in the attention computation. If cross-attention,
+ will be used for the query (i.e., key and value will use the encoder_hidden_states).
+ """
+ if input_ids is None and query_embeds is None:
+ raise ValueError("You have to specify query_embeds when input_ids is None")
+
+ query_length = query_embeds.shape[1] if query_embeds is not None else 0
+
+ embedding_output = self.embeddings(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ query_embeds=query_embeds,
+ )
+
+ input_shape = embedding_output.size()[:-1]
+ batch_size, seq_length = input_shape
+ device = embedding_output.device
+
+ if attention_mask is None:
+ attention_mask = torch.ones(((batch_size, seq_length)), device=device)
+
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+ # ourselves in which case we just need to make it broadcastable to all heads.
+ extended_attention_mask = self.get_extended_attention_mask(attention_mask, input_shape, device)
+
+ # If a 2D or 3D attention mask is provided for the cross-attention
+ # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+ if encoder_hidden_states is not None:
+ if isinstance(encoder_hidden_states, list):
+ encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states[0].size()
+ else:
+ encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+ encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+
+ if isinstance(encoder_attention_mask, list):
+ encoder_extended_attention_mask = [self.invert_attention_mask(mask) for mask in encoder_attention_mask]
+ elif encoder_attention_mask is None:
+ encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+ encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+ else:
+ encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+ else:
+ encoder_extended_attention_mask = None
+
+ encoder_outputs: BaseModelOutput = self.encoder(
+ embedding_output,
+ attention_mask=extended_attention_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_extended_attention_mask,
+ query_length=query_length,
+ **kwargs,
+ )
+ sequence_output = encoder_outputs.last_hidden_state
+ pooled_output = sequence_output[:, 0, :]
+
+ return BaseModelOutputWithPoolingAndCrossAttentions(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ InstructBLIP base Model consisting of language model, qformer and vision encoder.
+ """
+)
+class InstructBlipModel(InstructBlipPreTrainedModel):
+ main_input_name = "pixel_values"
+ _keep_in_fp32_modules = ["query_tokens"] # TODO @ArthurZucker I don't know why this is required for FP8
+
+ def __init__(self, config: InstructBlipConfig):
+ super().__init__(config)
+
+ self.vision_model = InstructBlipVisionModel(config.vision_config)
+ self.query_tokens = nn.Parameter(torch.zeros(1, config.num_query_tokens, config.qformer_config.hidden_size))
+ self.qformer = InstructBlipQFormerModel(config.qformer_config)
+
+ self.language_projection = nn.Linear(config.qformer_config.hidden_size, config.text_config.hidden_size)
+ self.language_model = AutoModel.from_config(config.text_config)
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.language_model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.language_model.set_input_embeddings(value)
+
+ def _preprocess_accelerate(self):
+ r"""
+ Some pre-processing hacks to make the model `accelerate` compatible. Check
+ https://github.com/huggingface/transformers/pull/21707 for more details.
+ """
+ hf_device_map = self.hf_device_map
+
+ if len(hf_device_map) > 1 and "language_model" not in hf_device_map and torch.cuda.device_count() > 1:
+ # warn users about unexpected behavior when using multi-GPU + InstructBLIP + `accelerate`.
+ logger.warning(
+ "The `language_model` is not in the `hf_device_map` dictionary and you are running your script"
+ " in a multi-GPU environment. this may lead to unexpected behavior when using `accelerate`."
+ " Please pass a `device_map` that contains `language_model` to remove this warning."
+ " Please refer to https://github.com/huggingface/blog/blob/main/accelerate-large-models.md for"
+ " more details on creating a `device_map` for large models.",
+ )
+
+ if hasattr(self.language_model, "_hf_hook"):
+ self.language_model._hf_hook.io_same_device = True # For `generate` compatibility
+
+ def get_placeholder_mask(self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ else:
+ special_image_mask = input_ids == self.config.image_token_id
+
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ return special_image_mask
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor,
+ qformer_input_ids: torch.FloatTensor,
+ qformer_attention_mask: torch.LongTensor | None = None,
+ input_ids: torch.FloatTensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ **kwargs: Unpack[FlashAttentionKwargs],
+ ) -> tuple | InstructBlipForConditionalGenerationModelOutput:
+ r"""
+ qformer_input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Indices of input sequence tokens in the vocabulary of the Q-Former. Input tokens can optionally be provided
+ to serve as text prompt, which the Q-Former model will encode.
+
+ Indices can be obtained using [`InstructBlipProcessor`]. See [`InstructBlipProcessor.__call__`] for
+ details.
+
+ [What are input IDs?](../glossary#input-ids)
+ qformer_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)
+ decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+ be used by default.
+
+ Only relevant in case an encoder-decoder language model (like T5) is used.
+ """
+
+ # step 1: forward the images through the vision encoder,
+ # to get image embeddings of shape (batch_size, seq_len, hidden_size)
+ vision_outputs = self.vision_model(
+ pixel_values=pixel_values,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ **kwargs,
+ )
+ image_embeds = vision_outputs[0]
+
+ # step 2: forward the query tokens through the QFormer, using the image embeddings for cross-attention
+ image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device)
+
+ # difference with BLIP-2 here: we also feed the instruction prompt to the Q-Former
+ query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)
+ query_attention_mask = torch.ones(query_tokens.size()[:-1], dtype=torch.long, device=image_embeds.device)
+ if qformer_attention_mask is None:
+ qformer_attention_mask = torch.ones_like(qformer_input_ids)
+ qformer_attention_mask = torch.cat([query_attention_mask, qformer_attention_mask], dim=1)
+ query_outputs = self.qformer(
+ input_ids=qformer_input_ids,
+ attention_mask=qformer_attention_mask,
+ query_embeds=query_tokens,
+ encoder_hidden_states=image_embeds,
+ encoder_attention_mask=image_attention_mask,
+ **kwargs,
+ )
+ query_output = query_outputs[0][:, : query_tokens.size(1), :]
+
+ if inputs_embeds is None:
+ inputs_embeds = self.language_model.get_input_embeddings()(input_ids)
+ if attention_mask is None:
+ attention_mask = torch.ones_like(input_ids)
+
+ # step 3: use the language model, conditioned on the query outputs and the prompt
+ language_model_inputs = self.language_projection(query_output)
+ language_model_inputs = language_model_inputs.to(inputs_embeds.device, inputs_embeds.dtype)
+ special_image_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds)
+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, language_model_inputs)
+
+ if self.config.use_decoder_only_language_model:
+ outputs = self.language_model(
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ else:
+ outputs = self.language_model(
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ decoder_input_ids=decoder_input_ids,
+ decoder_attention_mask=decoder_attention_mask,
+ **kwargs,
+ )
+
+ return InstructBlipForConditionalGenerationModelOutput(
+ vision_outputs=vision_outputs,
+ qformer_outputs=query_outputs,
+ language_model_outputs=outputs,
+ )
+
+
+@auto_docstring(
+ custom_intro="""
+ InstructBLIP Model for generating text given an image and an optional text prompt. The model consists of a vision
+ encoder, Querying Transformer (Q-Former) and a language model.
+
+ One can optionally pass `input_ids` to the model, which serve as a text prompt, to make the language model continue
+ the prompt. Otherwise, the language model starts generating text from the [BOS] (beginning-of-sequence) token.
+ """
+)
+class InstructBlipForConditionalGeneration(InstructBlipPreTrainedModel, GenerationMixin):
+ config: InstructBlipConfig
+ main_input_name = "pixel_values"
+
+ _can_compile_fullgraph = True
+ _keep_in_fp32_modules = ["query_tokens"] # TODO @ArthurZucker I don't know why this is required for FP8
+
+ def __init__(self, config: InstructBlipConfig):
+ super().__init__(config)
+
+ self.vision_model = InstructBlipVisionModel._from_config(config.vision_config)
+
+ self.query_tokens = nn.Parameter(torch.zeros(1, config.num_query_tokens, config.qformer_config.hidden_size))
+ self.qformer = InstructBlipQFormerModel._from_config(config.qformer_config)
+
+ self.language_projection = nn.Linear(config.qformer_config.hidden_size, config.text_config.hidden_size)
+
+ if config.use_decoder_only_language_model:
+ language_model = AutoModelForCausalLM.from_config(config.text_config)
+ else:
+ language_model = AutoModelForSeq2SeqLM.from_config(config.text_config)
+
+ self.language_model = language_model
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.language_model.get_input_embeddings()
+
+ def set_input_embeddings(self, value):
+ self.language_model.set_input_embeddings(value)
+
+ def set_output_embeddings(self, new_embeddings):
+ self.language_model.set_output_embeddings(new_embeddings)
+
+ def get_output_embeddings(self) -> nn.Module:
+ return self.language_model.get_output_embeddings()
+
+ def get_encoder(self, modality=None):
+ if modality is None:
+ return self.language_model.get_encoder()
+ else:
+ return super().get_encoder(modality=modality)
+
+ def get_decoder(self):
+ return self.language_model.get_decoder()
+
+ # Copied from transformers.models.instructblip.modeling_instructblip.InstructBlipModel._preprocess_accelerate
+ def _preprocess_accelerate(self):
+ r"""
+ Some pre-processing hacks to make the model `accelerate` compatible. Check
+ https://github.com/huggingface/transformers/pull/21707 for more details.
+ """
+ hf_device_map = self.hf_device_map
+
+ if len(hf_device_map) > 1 and "language_model" not in hf_device_map and torch.cuda.device_count() > 1:
+ # warn users about unexpected behavior when using multi-GPU + InstructBLIP + `accelerate`.
+ logger.warning(
+ "The `language_model` is not in the `hf_device_map` dictionary and you are running your script"
+ " in a multi-GPU environment. this may lead to unexpected behavior when using `accelerate`."
+ " Please pass a `device_map` that contains `language_model` to remove this warning."
+ " Please refer to https://github.com/huggingface/blog/blob/main/accelerate-large-models.md for"
+ " more details on creating a `device_map` for large models.",
+ )
+
+ if hasattr(self.language_model, "_hf_hook"):
+ self.language_model._hf_hook.io_same_device = True # For `generate` compatibility
+
+ @can_return_tuple
+ @auto_docstring
+ def get_image_features(
+ self,
+ pixel_values: torch.FloatTensor,
+ qformer_input_ids: torch.LongTensor,
+ qformer_attention_mask: torch.LongTensor | None = None,
+ interpolate_pos_encoding: bool | None = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | BaseModelOutputWithVisionQformerOutputs:
+ r"""
+ pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`):
+ The tensors corresponding to the input images.
+ qformer_input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Indices of input sequence tokens in the vocabulary of the Q-Former. Input tokens can optionally be provided
+ to serve as text prompt, which the Q-Former model will encode.
+
+ Indices can be obtained using [`InstructBlipProcessor`]. See [`InstructBlipProcessor.__call__`] for
+ details.
+
+ [What are input IDs?](../glossary#input-ids)
+ qformer_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)
+ """
+ # step 1: forward the images through the vision encoder,
+ # to get image embeddings of shape (batch_size, seq_len, hidden_size)
+ vision_outputs: BaseModelOutputWithPooling = self.vision_model(
+ pixel_values=pixel_values,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ return_dict=True,
+ **kwargs,
+ )
+ vision_outputs = BaseModelOutputWithVisionQformerOutputs(**vision_outputs, vision_outputs=vision_outputs)
+ image_embeds = vision_outputs[0]
+
+ # step 2: forward the query tokens through the QFormer, using the image embeddings for cross-attention
+ image_attention_mask = torch.ones(image_embeds.size()[:-1], dtype=torch.long, device=image_embeds.device)
+
+ # difference with BLIP-2 here: we also feed the instruction prompt to the Q-Former
+ query_tokens = self.query_tokens.expand(image_embeds.shape[0], -1, -1)
+ query_attention_mask = torch.ones(query_tokens.size()[:-1], dtype=torch.long, device=image_embeds.device)
+ if qformer_attention_mask is None:
+ qformer_attention_mask = torch.ones_like(qformer_input_ids)
+ qformer_attention_mask = torch.cat([query_attention_mask, qformer_attention_mask], dim=1)
+ qformer_outputs = self.qformer(
+ input_ids=qformer_input_ids,
+ attention_mask=qformer_attention_mask,
+ query_embeds=query_tokens,
+ encoder_hidden_states=image_embeds,
+ encoder_attention_mask=image_attention_mask,
+ return_dict=True,
+ **kwargs,
+ )
+ vision_outputs.qformer_outputs = qformer_outputs
+ query_output = qformer_outputs[0][:, : query_tokens.size(1), :]
+
+ # step 3: use the language model, conditioned on the query outputs and the prompt
+ image_features = self.language_projection(query_output)
+ vision_outputs.pooler_output = image_features
+
+ return vision_outputs
+
+ def get_placeholder_mask(self, input_ids: torch.LongTensor, inputs_embeds: torch.FloatTensor):
+ """
+ Obtains multimodal placeholder mask from `input_ids` or `inputs_embeds`.
+ """
+ if input_ids is None:
+ special_image_mask = inputs_embeds == self.get_input_embeddings()(
+ torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device)
+ )
+ special_image_mask = special_image_mask.all(-1)
+ else:
+ special_image_mask = input_ids == self.config.image_token_id
+
+ special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device)
+ return special_image_mask
+
+ @can_return_tuple
+ @auto_docstring
+ def forward(
+ self,
+ pixel_values: torch.FloatTensor,
+ qformer_input_ids: torch.FloatTensor,
+ qformer_attention_mask: torch.LongTensor | None = None,
+ input_ids: torch.FloatTensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ decoder_input_ids: torch.LongTensor | None = None,
+ decoder_attention_mask: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ labels: torch.LongTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ **kwargs: Unpack[TransformersKwargs],
+ ) -> tuple | InstructBlipForConditionalGenerationModelOutput:
+ r"""
+ qformer_input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Indices of input sequence tokens in the vocabulary of the Q-Former. Input tokens can optionally be provided
+ to serve as text prompt, which the Q-Former model will encode.
+
+ Indices can be obtained using [`InstructBlipProcessor`]. See [`InstructBlipProcessor.__call__`] for
+ details.
+
+ [What are input IDs?](../glossary#input-ids)
+ qformer_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)
+ decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+ Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+ be used by default.
+
+ Only relevant in case an encoder-decoder language model (like T5) is used.
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+ Labels for computing the language modeling loss. Indices should be in `[-100, 0, ..., config.vocab_size -
+ 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ...,
+ config.vocab_size]`
+
+ Examples:
+
+ ```python
+ >>> from transformers import InstructBlipProcessor, InstructBlipForConditionalGeneration
+ >>> import torch
+ >>> from PIL import Image
+ >>> import httpx
+ >>> from io import BytesIO
+
+ >>> model = InstructBlipForConditionalGeneration.from_pretrained("Salesforce/instructblip-vicuna-7b")
+ >>> processor = InstructBlipProcessor.from_pretrained("Salesforce/instructblip-vicuna-7b")
+
+ >>> device = "cuda" if torch.cuda.is_available() else "cpu"
+ >>> model.to(device) # doctest: +IGNORE_RESULT
+
+ >>> url = "https://raw.githubusercontent.com/salesforce/LAVIS/main/docs/_static/Confusing-Pictures.jpg"
+ >>> with httpx.stream("GET", url) as response:
+ ... image = Image.open(BytesIO(response.read())).convert("RGB")
+ >>> prompt = "What is unusual about this image?"
+ >>> inputs = processor(images=image, text=prompt, return_tensors="pt").to(device)
+
+ >>> outputs = model.generate(
+ ... **inputs,
+ ... do_sample=False,
+ ... num_beams=5,
+ ... max_length=256,
+ ... min_length=1,
+ ... top_p=0.9,
+ ... repetition_penalty=1.5,
+ ... length_penalty=1.0,
+ ... temperature=1,
+ ... )
+ >>> generated_text = processor.batch_decode(outputs, skip_special_tokens=True)[0].strip()
+ >>> print(generated_text)
+ The unusual aspect of this image is that a man is ironing clothes on the back of a yellow SUV, which is parked in the middle of a busy city street. This is an unconventional approach to ironing clothes, as it requires the man to balance himself and his ironing equipment on top of the vehicle while navigating through traffic. Additionally, the presence of taxis and other vehicles in the scene further emphasizes the unusual nature of this situation.
+ ```"""
+
+ image_features: BaseModelOutputWithVisionQformerOutputs = self.get_image_features(
+ pixel_values,
+ qformer_input_ids=qformer_input_ids,
+ qformer_attention_mask=qformer_attention_mask,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ return_dict=True,
+ )
+ language_model_inputs = image_features.pooler_output
+ qformer_outputs = image_features.qformer_outputs
+ vision_outputs = image_features.vision_outputs
+
+ if inputs_embeds is None:
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if attention_mask is None:
+ attention_mask = torch.ones_like(input_ids)
+
+ language_model_inputs = language_model_inputs.to(inputs_embeds.device, inputs_embeds.dtype)
+ special_image_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds)
+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, language_model_inputs)
+
+ if self.config.use_decoder_only_language_model:
+ outputs = self.language_model(
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+ logits = outputs[0]
+ loss = None
+ if labels is not None:
+ loss = self.loss_function(
+ logits=logits, labels=labels, vocab_size=self.config.text_config.vocab_size, **kwargs
+ )
+
+ else:
+ kwargs["return_dict"] = True
+ outputs = self.language_model(
+ inputs_embeds=inputs_embeds,
+ attention_mask=attention_mask,
+ decoder_input_ids=decoder_input_ids,
+ decoder_attention_mask=decoder_attention_mask,
+ labels=labels,
+ **kwargs,
+ )
+ loss = outputs.loss
+ logits = outputs.logits
+
+ return InstructBlipForConditionalGenerationModelOutput(
+ loss=loss,
+ logits=logits,
+ vision_outputs=vision_outputs,
+ qformer_outputs=qformer_outputs,
+ language_model_outputs=outputs,
+ )
+
+ @torch.no_grad()
+ def generate(
+ self,
+ pixel_values: torch.FloatTensor,
+ qformer_input_ids: torch.LongTensor | None = None,
+ qformer_attention_mask: torch.LongTensor | None = None,
+ input_ids: torch.LongTensor | None = None,
+ attention_mask: torch.LongTensor | None = None,
+ inputs_embeds: torch.FloatTensor | None = None,
+ interpolate_pos_encoding: bool = False,
+ **generate_kwargs,
+ ) -> torch.LongTensor:
+ """
+ Overrides `generate` function to be able to use the model as a conditional generator.
+
+ Args:
+ pixel_values (`torch.FloatTensor` of shape (batch_size, num_channels, height, width)):
+ Input images to be processed.
+ qformer_input_ids (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*):
+ The sequence used as a prompt to be fed to the Q-Former module.
+ qformer_attention_mask (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*):
+ Mask to avoid performing attention on padding token indices.
+ input_ids (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*):
+ The sequence used as a prompt for the generation.
+ attention_mask (`torch.LongTensor` of shape (batch_size, sequence_length), *optional*):
+ Mask to avoid performing attention on padding token indices.
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):
+ Embedded representation of the inputs. Should be float, not int tokens.
+ interpolate_pos_encoding (`bool`, *optional*, defaults to `False`):
+ Whether to interpolate the positional encoding of the image embeddings.
+
+ Returns:
+ captions (list): A list of strings of length batch_size * num_captions.
+ """
+ if hasattr(self, "hf_device_map"):
+ # preprocess for `accelerate`
+ self._preprocess_accelerate()
+
+ batch_size = pixel_values.shape[0]
+ image_features: BaseModelOutputWithVisionQformerOutputs = self.get_image_features(
+ pixel_values,
+ qformer_input_ids=qformer_input_ids,
+ qformer_attention_mask=qformer_attention_mask,
+ interpolate_pos_encoding=interpolate_pos_encoding,
+ return_dict=True,
+ )
+ language_model_inputs = image_features.pooler_output
+
+ if inputs_embeds is None:
+ if input_ids is None:
+ image_tokens = [self.config.image_token_index] * self.config.num_query_tokens
+ start_tokens = image_tokens + [self.config.text_config.bos_token_id]
+ input_ids = torch.tensor([start_tokens], dtype=torch.long, device=pixel_values.device)
+ input_ids = input_ids.repeat(batch_size, 1)
+ inputs_embeds = self.get_input_embeddings()(input_ids)
+
+ if attention_mask is None:
+ attention_mask = torch.ones_like(input_ids)
+
+ language_model_inputs = language_model_inputs.to(inputs_embeds.device, inputs_embeds.dtype)
+ special_image_mask = self.get_placeholder_mask(input_ids, inputs_embeds=inputs_embeds)
+ inputs_embeds = inputs_embeds.masked_scatter(special_image_mask, language_model_inputs)
+
+ inputs = {"inputs_embeds": inputs_embeds, "attention_mask": attention_mask}
+ if not self.language_model.config.is_encoder_decoder:
+ inputs["input_ids"] = input_ids
+
+ outputs = self.language_model.generate(**inputs, **generate_kwargs)
+
+ return outputs
+
+
+__all__ = [
+ "InstructBlipQFormerModel",
+ "InstructBlipPreTrainedModel",
+ "InstructBlipModel",
+ "InstructBlipForConditionalGeneration",
+ "InstructBlipVisionModel",
+]
diff --git a/third_party/transformers/src/transformers/models/instructblip/processing_instructblip.py b/third_party/transformers/src/transformers/models/instructblip/processing_instructblip.py
new file mode 100644
index 0000000000000000000000000000000000000000..10c69dd79d1d4d7b2a65a6f326dbad9a3f40b3be
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/instructblip/processing_instructblip.py
@@ -0,0 +1,123 @@
+# Copyright 2023 The HuggingFace Inc. team.
+#
+# 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.
+"""
+Processor class for InstructBLIP. Largely copy of Blip2Processor with addition of a tokenizer for the Q-Former.
+"""
+
+from ...image_processing_utils import BatchFeature
+from ...image_utils import ImageInput
+from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack
+from ...tokenization_utils_base import AddedToken, PreTokenizedInput, TextInput
+from ...utils import auto_docstring, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class InstructBlipProcessorKwargs(ProcessingKwargs, total=False):
+ _defaults = {
+ "text_kwargs": {
+ "add_special_tokens": True,
+ "padding": False,
+ "stride": 0,
+ "return_overflowing_tokens": False,
+ "return_special_tokens_mask": False,
+ "return_offsets_mapping": False,
+ "return_token_type_ids": False,
+ "return_length": False,
+ "verbose": True,
+ },
+ }
+
+
+@auto_docstring
+class InstructBlipProcessor(ProcessorMixin):
+ def __init__(self, image_processor, tokenizer, qformer_tokenizer, num_query_tokens=None, **kwargs):
+ r"""
+ qformer_tokenizer (`AutoTokenizer`):
+ An instance of ['PreTrainedTokenizer`]. The Q-Former tokenizer is a required input.
+ num_query_tokens (`int`, *optional*):
+ "
+ Number of tokens used by the Qformer as queries, should be same as in model's config.
+ """
+ if not hasattr(tokenizer, "image_token"):
+ self.image_token = AddedToken("", normalized=False, special=True)
+ tokenizer.add_tokens([self.image_token], special_tokens=True)
+ else:
+ self.image_token = tokenizer.image_token
+ self.num_query_tokens = num_query_tokens
+
+ super().__init__(image_processor, tokenizer, qformer_tokenizer)
+
+ @auto_docstring
+ def __call__(
+ self,
+ images: ImageInput | None = None,
+ text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None,
+ **kwargs: Unpack[InstructBlipProcessorKwargs],
+ ) -> BatchFeature:
+ if images is None and text is None:
+ raise ValueError("You have to specify at least images or text.")
+
+ output_kwargs = self._merge_kwargs(
+ InstructBlipProcessorKwargs,
+ tokenizer_init_kwargs=self.tokenizer.init_kwargs,
+ **kwargs,
+ )
+
+ return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None)
+ encoding = {}
+ if text is not None:
+ if isinstance(text, str):
+ text = [text]
+ elif not isinstance(text, list) and not isinstance(text[0], str):
+ raise ValueError("Invalid input text. Please provide a string, or a list of strings")
+
+ qformer_text_encoding = self.qformer_tokenizer(text, **output_kwargs["text_kwargs"])
+ encoding["qformer_input_ids"] = qformer_text_encoding.pop("input_ids")
+ encoding["qformer_attention_mask"] = qformer_text_encoding.pop("attention_mask")
+
+ # We need this hacky manipulation because BLIP expects image tokens to be at the beginning even before BOS token
+ if output_kwargs["text_kwargs"].get("max_length") is not None:
+ output_kwargs["text_kwargs"]["max_length"] -= self.num_query_tokens
+ text_encoding = self.tokenizer(text, **output_kwargs["text_kwargs"])
+
+ if images is not None:
+ # Image tokens should not be padded/truncated or prepended with special BOS token
+ image_tokens = self.image_token.content * self.num_query_tokens
+ output_kwargs["text_kwargs"]["add_special_tokens"] = False
+ output_kwargs["text_kwargs"]["padding"] = False
+ output_kwargs["text_kwargs"]["truncation"] = False
+ image_text_encoding = self.tokenizer(image_tokens, **output_kwargs["text_kwargs"])
+ for k in text_encoding:
+ text_encoding[k] = [image_text_encoding[k] + sample for sample in text_encoding[k]]
+ encoding.update(text_encoding)
+
+ if images is not None:
+ image_encoding = self.image_processor(images, **output_kwargs["images_kwargs"])
+ encoding.update(image_encoding)
+
+ # Cast to desired return tensors type
+ encoding = BatchFeature(encoding, tensor_type=return_tensors)
+ return encoding
+
+ @property
+ def model_input_names(self):
+ tokenizer_input_names = self.tokenizer.model_input_names
+ image_processor_input_names = self.image_processor.model_input_names
+ qformer_input_names = ["qformer_input_ids", "qformer_attention_mask"]
+ return tokenizer_input_names + image_processor_input_names + qformer_input_names
+
+
+__all__ = ["InstructBlipProcessor"]
diff --git a/third_party/transformers/src/transformers/models/internvl/__init__.py b/third_party/transformers/src/transformers/models/internvl/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6d4ffe7befafddc423a58d41934a260137da0a13
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/internvl/__init__.py
@@ -0,0 +1,29 @@
+# Copyright 2025 The HuggingFace Team. All rights reserved.
+#
+# 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.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+ from .configuration_internvl import *
+ from .modeling_internvl import *
+ from .processing_internvl import *
+ from .video_processing_internvl import *
+else:
+ import sys
+
+ _file = globals()["__file__"]
+ sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)
diff --git a/third_party/transformers/src/transformers/models/internvl/configuration_internvl.py b/third_party/transformers/src/transformers/models/internvl/configuration_internvl.py
new file mode 100644
index 0000000000000000000000000000000000000000..e6a4834102a20a5af7bd3b77fca66ac489409105
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/internvl/configuration_internvl.py
@@ -0,0 +1,136 @@
+# Copyright 2025 HuggingFace Inc. team. All rights reserved.
+#
+# 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.
+
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring
+from ..auto import CONFIG_MAPPING, AutoConfig
+
+
+@auto_docstring(checkpoint="OpenGVLab/InternVL3-1B-hf")
+@strict
+class InternVLVisionConfig(PreTrainedConfig):
+ r"""
+ projection_dropout (`float`, *optional*, defaults to 0.0):
+ Dropout probability for the projection layer.
+ norm_type (`str`, *optional*, defaults to `"layer_norm"`):
+ The type of normalization to use in the encoder. Can be `"layer_norm"` or `"rms_norm"`.
+ use_mask_token (`bool`, *optional*, defaults to `False`):
+ Whether to use a mask token for masked image modeling
+ use_mean_pooling (`bool`, *optional*, defaults to `True`):
+ Whether to mean pool the final hidden states of the patches instead of using the final hidden state of the
+ CLS token, before applying the classification head.
+
+ Example:
+
+ ```python
+ >>> from transformers import InternVLVisionConfig, InternVLVisionModel
+
+ >>> # Initializing a InternVLVisionModel OpenGVLab/InternVL3-1B-hf style configuration
+ >>> configuration = InternVLVisionConfig()
+
+ >>> # Initializing a model (with random weights) from the OpenGVLab/InternVL3-1B-hf configuration
+ >>> model = InternVLVisionModel(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "internvl_vision"
+ base_config_key = "vision_config"
+
+ hidden_size: int = 1024
+ num_hidden_layers: int = 24
+ num_attention_heads: int = 16
+ attention_bias: bool = False
+ use_qk_norm: bool = False
+ intermediate_size: int = 4096
+ hidden_act: str = "gelu"
+ hidden_dropout_prob: float | int = 0.0
+ attention_dropout: float | int = 0.0
+ projection_dropout: float | int = 0.0
+ initializer_range: float = 0.02
+ norm_type: str = "layer_norm"
+ layer_norm_eps: float = 1e-06
+ image_size: int | list[int] | tuple[int, ...] = (448, 448)
+ patch_size: int | list[int] | tuple[int, ...] = (14, 14)
+ num_channels: int = 3
+ use_mask_token: bool = False
+ use_absolute_position_embeddings: bool = True
+ layer_scale_init_value: float = 0.1
+ use_mean_pooling: bool = True
+
+ def __post_init__(self, **kwargs):
+ self.image_size = (
+ self.image_size if isinstance(self.image_size, (list, tuple)) else (self.image_size, self.image_size)
+ )
+ self.patch_size = (
+ self.patch_size if isinstance(self.patch_size, (list, tuple)) else (self.patch_size, self.patch_size)
+ )
+ super().__post_init__(**kwargs)
+
+
+@auto_docstring(checkpoint="OpenGVLab/InternVL3-1B-hf")
+@strict
+class InternVLConfig(PreTrainedConfig):
+ r"""
+ downsample_ratio (`float`, *optional*, defaults to 0.5):
+ Factor by which to downsample the image.
+
+ Example:
+
+ ```python
+ >>> from transformers import InternVLForConditionalGeneration, InternVLConfig
+
+ >>> # Initializing a InternVL style configuration
+ >>> configuration = InternVLConfig()
+
+ >>> # Initializing a model (with random weights) from the OpenGVLab/InternVL3-1B-hf configuration
+ >>> model = InternVLForConditionalGeneration(configuration)
+
+ >>> # Accessing the model configuration
+ >>> configuration = model.config
+ ```"""
+
+ model_type = "internvl"
+ sub_configs = {"text_config": AutoConfig, "vision_config": InternVLVisionConfig}
+
+ vision_config: dict | PreTrainedConfig | None = None
+ text_config: dict | PreTrainedConfig | None = None
+ image_token_id: int = 151667
+ image_seq_length: int = 256
+ downsample_ratio: float = 0.5
+ projector_hidden_act: str = "gelu"
+ vision_feature_layer: int | list[int] = -1
+ vision_feature_select_strategy: str = "default"
+ tie_word_embeddings: bool = True
+
+ def __post_init__(self, **kwargs):
+ if isinstance(self.vision_config, dict):
+ self.vision_config = InternVLVisionConfig(**self.vision_config)
+ elif self.vision_config is None:
+ self.vision_config = InternVLVisionConfig()
+
+ if isinstance(self.text_config, dict):
+ self.text_config["model_type"] = self.text_config.get("model_type", "qwen2")
+ self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config)
+ elif self.text_config is None:
+ self.text_config = CONFIG_MAPPING["qwen2"]()
+
+ super().__post_init__(**kwargs)
+
+
+__all__ = ["InternVLVisionConfig", "InternVLConfig"]
diff --git a/third_party/transformers/src/transformers/models/internvl/convert_internvl_weights_to_hf.py b/third_party/transformers/src/transformers/models/internvl/convert_internvl_weights_to_hf.py
new file mode 100644
index 0000000000000000000000000000000000000000..e6892ded6b804570ca37dcea2d174767b66513cb
--- /dev/null
+++ b/third_party/transformers/src/transformers/models/internvl/convert_internvl_weights_to_hf.py
@@ -0,0 +1,463 @@
+# Copyright 2025 HuggingFace Inc. team. All rights reserved.
+# 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.
+import argparse
+import gc
+import os
+import re
+from typing import Literal
+
+import torch
+from einops import rearrange
+
+from transformers import (
+ AutoModel,
+ AutoTokenizer,
+ GenerationConfig,
+ GotOcr2ImageProcessor,
+ InternVLConfig,
+ InternVLForConditionalGeneration,
+ InternVLProcessor,
+ InternVLVideoProcessor,
+ InternVLVisionConfig,
+ LlamaConfig,
+ Qwen2Config,
+)
+
+
+LM_TYPE_CORRESPONDENCE = {
+ "OpenGVLab/InternVL2_5-1B-MPO": "qwen2",
+ "OpenGVLab/InternVL2_5-2B-MPO": "llama",
+ "OpenGVLab/InternVL2_5-4B-MPO": "qwen2",
+ "OpenGVLab/InternVL2_5-8B-MPO": "llama",
+ "OpenGVLab/InternVL2_5-26B-MPO": "llama",
+ "OpenGVLab/InternVL2_5-38B-MPO": "qwen2",
+ "OpenGVLab/InternVL2_5-78B-MPO": "qwen2",
+ "OpenGVLab/InternVL3-1B": "qwen2",
+ "OpenGVLab/InternVL3-2B": "qwen2",
+ "OpenGVLab/InternVL3-8B": "qwen2",
+ "OpenGVLab/InternVL3-9B": "llama",
+ "OpenGVLab/InternVL3-14B": "qwen2",
+ "OpenGVLab/InternVL3-38B": "qwen2",
+ "OpenGVLab/InternVL3-78B": "qwen2",
+}
+
+UNNECESSARY_CONFIG_KEYS = [ "_name_or_path", "_attn_implementation_autoset", "auto_map", "use_bfloat16", "use_flash_attn", "bias", "laux_allreduce", "moe_coeff_ratio", "moe_intermediate_size", "moe_output_scale", "noisy_gate_policy", "shared_expert_intermediate_size", "use_residual", "use_moe", "use_rts", "use_weighted_residual", "moe_config", "num_experts", "num_routed_experts", "num_shared_experts", "capacity_factor", "eval_capacity_factor", "drop_path_rate"] # fmt: skip
+
+# fmt: off
+ORIGINAL_TO_CONVERTED_KEY_MAPPING_VISION = {
+ # Vision encoder mapping
+ r"vision_model": r"model.vision_tower",
+ r"layers": r"layer",
+ r"class_embedding": r"cls_token",
+ r"position_embedding": r"position_embeddings",
+ r"patch_embedding": r"patch_embeddings.projection",
+ r"ls(\d+)": r"lambda_\1",
+ r"attn.proj": r"attention.projection_layer",
+ r"attn.dropout": r"attention.projection_dropout",
+ r"attn": r"attention",
+ r"norm1": r"layernorm_before",
+ r"norm2": r"layernorm_after",
+
+}
+
+ORIGINAL_TO_CONVERTED_KEY_MAPPING_TEXT_LLAMA = {
+ r"language_model.model.": r"model.language_model.",
+ r"tok_embeddings": r"embed_tokens",
+ r"attention.wo": r"self_attn.o_proj",
+ r"feed_forward.w1": r"mlp.gate_proj",
+ r"feed_forward.w2": r"mlp.down_proj",
+ r"feed_forward.w3": r"mlp.up_proj",
+ r"attention_norm": r"input_layernorm",
+ r"ffn_norm": r"post_attention_layernorm",
+ r"language_model.output": r"lm_head",
+}
+
+ORIGINAL_TO_CONVERTED_KEY_MAPPING_TEXT_QWEN2 = {
+ # Vision encoder mapping
+ r"language_model.model.": r"model.language_model.",
+ r"language_model.lm_head": r"lm_head",
+}
+
+ORIGINAL_TO_CONVERTED_KEY_MAPPING_MULTI = {
+ # Vision encoder mapping
+ r"mlp1.0": r"model.multi_modal_projector.layer_norm",
+ r"mlp1.1": r"model.multi_modal_projector.linear_1",
+ r"mlp1.3": r"model.multi_modal_projector.linear_2",
+}
+
+
+chat_template = (
+ "{% for message in messages %}"
+ "{{'<|im_start|>' + message['role'] + '\n'}}"
+ "{% if message['content'] is string %}"
+ "{{ message['content'] }}"
+ "{% else %}"
+ "{% for content in message['content'] %}"
+ "{% if content['type'] == 'image' %}"
+ "{{ '\n' }}"
+ "{% elif content['type'] == 'video' %}"
+ "{{ '