diff --git a/third_party/transformers/src/transformers/models/cohere2/modeling_cohere2.py b/third_party/transformers/src/transformers/models/cohere2/modeling_cohere2.py new file mode 100644 index 0000000000000000000000000000000000000000..f43b2a0ef4122732259cd2e02e38ccff3bdbb478 --- /dev/null +++ b/third_party/transformers/src/transformers/models/cohere2/modeling_cohere2.py @@ -0,0 +1,509 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/cohere2/modular_cohere2.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_cohere2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2024 Cohere Inc. 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_kernelized_func +from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask +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 maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_cohere2 import Cohere2Config + + +class Cohere2RotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: Cohere2Config, 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: Cohere2Config | 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) + 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.repeat_interleave(freqs, 2, dim=-1) # diff from Llama: we interleave() instead of cat() + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class Cohere2LayerNorm(nn.Module): + def __init__(self, hidden_size=None, eps=1e-5, bias=False): + """The hidden size can be a tuple or an int. The tuple is used for QKNorm to normalize across head_dim""" + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + mean = hidden_states.mean(-1, keepdim=True) + variance = (hidden_states - mean).pow(2).mean(-1, keepdim=True) + hidden_states = (hidden_states - mean) * torch.rsqrt(variance + self.variance_epsilon) + hidden_states = self.weight.to(torch.float32) * hidden_states + return hidden_states.to(input_dtype) + + +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): + # Split and rotate. Note that this function is different from e.g. Llama. + x1 = x[..., ::2] + x2 = x[..., 1::2] + rot_x = torch.stack([-x2, x1], dim=-1).flatten(-2) + return rot_x + + +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. + """ + dtype = q.dtype + q = q.float() + k = k.float() + 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.to(dtype=dtype), k_embed.to(dtype=dtype) + + +@use_kernelized_func(apply_rotary_pos_emb) +class Cohere2Attention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Cohere2Config, 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 + layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None + self.sliding_window = config.sliding_window if layer_type == "sliding_attention" else None + + 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], + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> 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).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 + if self.sliding_window is not None: + 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, + sliding_window=self.sliding_window, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Cohere2MLP(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x): + down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + return down_proj + + +class Cohere2DecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Cohere2Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.self_attn = Cohere2Attention(config=config, layer_idx=layer_idx) + self.mlp = Cohere2MLP(config) + self.input_layernorm = Cohere2LayerNorm(hidden_size=(config.hidden_size), eps=config.layer_norm_eps) + + 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, + use_cache: bool | None = False, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + """ + 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_size, sequence_length)` if flash attention is used or `(batch_size, 1, + query_sequence_length, key_sequence_length)` if default attention is used. + past_key_values (`Cache`, *optional*): cached past key and value projection states + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + use_cache (`bool`, *optional*): + If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding + (see `past_key_values`). + 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. + """ + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states_attention, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states_mlp = self.mlp(hidden_states) + hidden_states = residual + hidden_states_attention + hidden_states_mlp + return hidden_states + + +@auto_docstring +class Cohere2PreTrainedModel(PreTrainedModel): + config: Cohere2Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Cohere2DecoderLayer"] + _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": Cohere2DecoderLayer, + "attentions": Cohere2Attention, + } + + +@auto_docstring +class Cohere2Model(Cohere2PreTrainedModel): + def __init__(self, config: Cohere2Config): + 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( + [Cohere2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = Cohere2LayerNorm(hidden_size=(config.hidden_size), eps=config.layer_norm_eps) + self.rotary_emb = Cohere2RotaryEmbedding(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 = 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) + + if not isinstance(causal_mask_mapping := attention_mask, dict): + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + causal_mask_mapping = { + "full_attention": create_causal_mask(**mask_kwargs), + "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs), + } + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for i, decoder_layer in enumerate(self.layers): + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask_mapping[self.config.layer_types[i]], + position_embeddings=position_embeddings, + past_key_values=past_key_values, + use_cache=use_cache, + position_ids=position_ids, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class Cohere2ForCausalLM(Cohere2PreTrainedModel, 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 = Cohere2Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + self.logit_scale = config.logit_scale + self.tie_word_embeddings = config.tie_word_embeddings + + # 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, Cohere2ForCausalLM + + >> model = Cohere2ForCausalLM.from_pretrained("Cohere2ForAI/c4ai-command-r-v01") + >> tokenizer = AutoTokenizer.from_pretrained("Cohere2ForAI/c4ai-command-r-v01") + + >> 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 + 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.logit_scale # main diff from Llama + + 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__ = ["Cohere2ForCausalLM", "Cohere2Model", "Cohere2PreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/cohere2/modular_cohere2.py b/third_party/transformers/src/transformers/models/cohere2/modular_cohere2.py new file mode 100644 index 0000000000000000000000000000000000000000..d19055a1b787fa0af17d710da42c2c7c4ac55ac1 --- /dev/null +++ b/third_party/transformers/src/transformers/models/cohere2/modular_cohere2.py @@ -0,0 +1,325 @@ +# Copyright 2024 Cohere Inc. 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 + +import torch +import torch.nn as nn +from huggingface_hub.dataclasses import strict + +from ...cache_utils import Cache, DynamicCache +from ...configuration_utils import PreTrainedConfig +from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask +from ...modeling_outputs import BaseModelOutputWithPast +from ...modeling_rope_utils import ( + RopeParameters, + dynamic_rope_update, +) +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, logging +from ...utils.generic import maybe_autocast +from ..cohere.modeling_cohere import ( + CohereAttention, + CohereDecoderLayer, + CohereForCausalLM, + CohereLayerNorm, + CoherePreTrainedModel, + CohereRotaryEmbedding, + apply_rotary_pos_emb, + eager_attention_forward, +) +from ..gemma2.modeling_gemma2 import Gemma2Model + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="CohereForAI/c4ai-command-r-v01") +@strict +class Cohere2Config(PreTrainedConfig): + r""" + logit_scale (`float`, *optional*, defaults to 0.0625): + The scaling factor for the output logits. + + ```python + >>> from transformers import Cohere2Model, Cohere2Config + + >>> # Initializing a Cohere Nextmodel configuration + >>> configuration = Cohere2Config() + + >>> # Initializing a model from the Cohere2 configuration + >>> model = Cohere2Model(configuration) # doctest: +SKIP + + >>> # Accessing the model configuration + >>> configuration = model.config # doctest: +SKIP + ``` + """ + + model_type = "cohere2" + 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_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"]), + } + + vocab_size: int = 256000 + hidden_size: int = 8192 + intermediate_size: int = 22528 + logit_scale: float = 0.0625 + num_hidden_layers: int = 40 + num_attention_heads: int = 64 + num_key_value_heads: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 8192 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-5 + use_cache: bool = True + pad_token_id: int | None = 0 + bos_token_id: int | None = 5 + eos_token_id: int | list[int] | None = 255001 + tie_word_embeddings: bool = True + rope_parameters: RopeParameters | dict | None = None + attention_bias: bool = False + attention_dropout: float | int = 0.0 + sliding_window: int | None = 4096 + layer_types: list[str] | None = None + + def __post_init__(self, **kwargs): + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + # Need to specify head_dim in the config so it can be used in the attention forward functions + self.head_dim = self.hidden_size // self.num_attention_heads + + # BC -> the pattern used to be a simple int, and it's still present in configs on the Hub + if self.layer_types is None: + # BC -> the pattern used to be a simple int, and it's still present in configs on the Hub + _sliding_window_pattern = kwargs.pop("sliding_window_pattern", 4) + self.layer_types = [ + "sliding_attention" if bool((i + 1) % _sliding_window_pattern) else "full_attention" + for i in range(self.num_hidden_layers) + ] + + super().__post_init__(**kwargs) + + +class Cohere2RotaryEmbedding(CohereRotaryEmbedding): + @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) + 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.repeat_interleave(freqs, 2, dim=-1) # diff from Llama: we interleave() instead of cat() + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class Cohere2LayerNorm(CohereLayerNorm): + pass + + +class Cohere2Attention(CohereAttention): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Cohere2Config, layer_idx: int | None = None): + nn.Module.__init__(self) + 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 + layer_type = config.layer_types[layer_idx] if hasattr(config, "layer_types") else None + self.sliding_window = config.sliding_window if layer_type == "sliding_attention" else None + + 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], + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> 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).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 + if self.sliding_window is not None: + 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, + sliding_window=self.sliding_window, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Cohere2DecoderLayer(CohereDecoderLayer): + def __init__(self, config: Cohere2Config, layer_idx: int): + super().__init__(config, layer_idx) + + 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, + use_cache: bool | None = False, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.FloatTensor, tuple[torch.FloatTensor, torch.FloatTensor] | None]: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states_attention, _ = self.self_attn( + hidden_states=hidden_states, + position_embeddings=position_embeddings, + attention_mask=attention_mask, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states_mlp = self.mlp(hidden_states) + hidden_states = residual + hidden_states_attention + hidden_states_mlp + return hidden_states + + +class Cohere2PreTrainedModel(CoherePreTrainedModel): + config: Cohere2Config + _can_record_outputs = { + "hidden_states": Cohere2DecoderLayer, + "attentions": Cohere2Attention, + } + + +class Cohere2Model(Gemma2Model): + def __init__(self, config: Cohere2Config): + super().__init__(config) + self.norm = Cohere2LayerNorm(hidden_size=(config.hidden_size), eps=config.layer_norm_eps) + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + + 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 = 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) + + if not isinstance(causal_mask_mapping := attention_mask, dict): + mask_kwargs = { + "config": self.config, + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "past_key_values": past_key_values, + "position_ids": position_ids, + } + causal_mask_mapping = { + "full_attention": create_causal_mask(**mask_kwargs), + "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs), + } + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for i, decoder_layer in enumerate(self.layers): + hidden_states = decoder_layer( + hidden_states, + attention_mask=causal_mask_mapping[self.config.layer_types[i]], + position_embeddings=position_embeddings, + past_key_values=past_key_values, + use_cache=use_cache, + position_ids=position_ids, + **kwargs, + ) + + hidden_states = self.norm(hidden_states) + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +class Cohere2ForCausalLM(CohereForCausalLM): + pass + + +__all__ = ["Cohere2Config", "Cohere2ForCausalLM", "Cohere2Model", "Cohere2PreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/conditional_detr/__init__.py b/third_party/transformers/src/transformers/models/conditional_detr/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..98371606d8eedd6d4c930228300bacf9c808be79 --- /dev/null +++ b/third_party/transformers/src/transformers/models/conditional_detr/__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_conditional_detr import * + from .image_processing_conditional_detr import * + from .image_processing_pil_conditional_detr import * + from .modeling_conditional_detr 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/conditional_detr/configuration_conditional_detr.py b/third_party/transformers/src/transformers/models/conditional_detr/configuration_conditional_detr.py new file mode 100644 index 0000000000000000000000000000000000000000..12f7977d7423d532f32aef41dd197dadb22a117c --- /dev/null +++ b/third_party/transformers/src/transformers/models/conditional_detr/configuration_conditional_detr.py @@ -0,0 +1,119 @@ +# 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. +"""Conditional DETR model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...backbone_utils import consolidate_backbone_kwargs_to_config +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring +from ..auto import AutoConfig + + +@auto_docstring(checkpoint="microsoft/conditional-detr-resnet-50") +@strict +class ConditionalDetrConfig(PreTrainedConfig): + r""" + num_queries (`int`, *optional*, defaults to 100): + Number of object queries, i.e. detection slots. This is the maximal number of objects + [`ConditionalDetrModel`] can detect in a single image. For COCO, we recommend 100 queries. + auxiliary_loss (`bool`, *optional*, defaults to `False`): + Whether auxiliary decoding losses (loss at each decoder layer) are to be used. + position_embedding_type (`str`, *optional*, defaults to `"sine"`): + Type of position embeddings to be used on top of the image features. One of `"sine"` or `"learned"`. + dilation (`bool`, *optional*, defaults to `False`): + Whether to replace stride with dilation in the last convolutional block (DC5). Only supported when + `use_timm_backbone` = `True`. + + Examples: + + ```python + >>> from transformers import ConditionalDetrConfig, ConditionalDetrModel + + >>> # Initializing a Conditional DETR microsoft/conditional-detr-resnet-50 style configuration + >>> configuration = ConditionalDetrConfig() + + >>> # Initializing a model (with random weights) from the microsoft/conditional-detr-resnet-50 style configuration + >>> model = ConditionalDetrModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "conditional_detr" + sub_configs = {"backbone_config": AutoConfig} + keys_to_ignore_at_inference = ["past_key_values"] + attribute_map = { + "hidden_size": "d_model", + "num_attention_heads": "encoder_attention_heads", + "num_hidden_layers": "encoder_layers", + } + + backbone_config: dict | PreTrainedConfig | None = None + num_channels: int = 3 + num_queries: int = 300 + encoder_layers: int = 6 + encoder_ffn_dim: int = 2048 + encoder_attention_heads: int = 8 + decoder_layers: int = 6 + decoder_ffn_dim: int = 2048 + decoder_attention_heads: int = 8 + encoder_layerdrop: float | int = 0.0 + decoder_layerdrop: float | int = 0.0 + is_encoder_decoder: bool = True + activation_function: str = "relu" + d_model: int = 256 + dropout: float | int = 0.1 + attention_dropout: float | int = 0.0 + activation_dropout: float | int = 0.0 + init_std: float = 0.02 + init_xavier_std: float = 1.0 + auxiliary_loss: bool = False + position_embedding_type: str = "sine" + dilation: bool = False + class_cost: int = 2 + bbox_cost: int = 5 + giou_cost: int = 2 + mask_loss_coefficient: int = 1 + dice_loss_coefficient: int = 1 + cls_loss_coefficient: int = 2 + bbox_loss_coefficient: int = 5 + giou_loss_coefficient: int = 2 + focal_alpha: float = 0.25 + + def __post_init__(self, **kwargs): + # Init timm backbone with hardcoded values for BC + backbone_kwargs = kwargs.get("backbone_kwargs", {}) + timm_default_kwargs = { + "num_channels": backbone_kwargs.get("num_channels", self.num_channels), + "features_only": True, + "use_pretrained_backbone": False, + "out_indices": backbone_kwargs.get("out_indices", [1, 2, 3, 4]), + } + if self.dilation: + timm_default_kwargs["output_stride"] = backbone_kwargs.get("output_stride", 16) + + self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config( + backbone_config=self.backbone_config, + default_backbone="resnet50", + default_config_type="resnet", + default_config_kwargs={"out_features": ["stage4"]}, + timm_default_kwargs=timm_default_kwargs, + **kwargs, + ) + + super().__post_init__(**kwargs) + + +__all__ = ["ConditionalDetrConfig"] diff --git a/third_party/transformers/src/transformers/models/conditional_detr/convert_conditional_detr_original_pytorch_checkpoint_to_pytorch.py b/third_party/transformers/src/transformers/models/conditional_detr/convert_conditional_detr_original_pytorch_checkpoint_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..fad35cb231ba486451154a935790c400baca9dd5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/conditional_detr/convert_conditional_detr_original_pytorch_checkpoint_to_pytorch.py @@ -0,0 +1,363 @@ +# 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 Conditional DETR checkpoints.""" + +import argparse +import json +from collections import OrderedDict +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 ( + ConditionalDetrConfig, + ConditionalDetrForObjectDetection, + ConditionalDetrForSegmentation, + ConditionalDetrImageProcessor, +) +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) +rename_keys = [] +for i in range(6): + # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms + rename_keys.append( + (f"transformer.encoder.layers.{i}.self_attn.out_proj.weight", f"encoder.layers.{i}.self_attn.out_proj.weight") + ) + rename_keys.append( + (f"transformer.encoder.layers.{i}.self_attn.out_proj.bias", f"encoder.layers.{i}.self_attn.out_proj.bias") + ) + rename_keys.append((f"transformer.encoder.layers.{i}.linear1.weight", f"encoder.layers.{i}.fc1.weight")) + rename_keys.append((f"transformer.encoder.layers.{i}.linear1.bias", f"encoder.layers.{i}.fc1.bias")) + rename_keys.append((f"transformer.encoder.layers.{i}.linear2.weight", f"encoder.layers.{i}.fc2.weight")) + rename_keys.append((f"transformer.encoder.layers.{i}.linear2.bias", f"encoder.layers.{i}.fc2.bias")) + rename_keys.append( + (f"transformer.encoder.layers.{i}.norm1.weight", f"encoder.layers.{i}.self_attn_layer_norm.weight") + ) + rename_keys.append((f"transformer.encoder.layers.{i}.norm1.bias", f"encoder.layers.{i}.self_attn_layer_norm.bias")) + rename_keys.append((f"transformer.encoder.layers.{i}.norm2.weight", f"encoder.layers.{i}.final_layer_norm.weight")) + rename_keys.append((f"transformer.encoder.layers.{i}.norm2.bias", f"encoder.layers.{i}.final_layer_norm.bias")) + # decoder layers: 2 times output projection, 2 feedforward neural networks and 3 layernorms + rename_keys.append( + (f"transformer.decoder.layers.{i}.self_attn.out_proj.weight", f"decoder.layers.{i}.self_attn.out_proj.weight") + ) + rename_keys.append( + (f"transformer.decoder.layers.{i}.self_attn.out_proj.bias", f"decoder.layers.{i}.self_attn.out_proj.bias") + ) + rename_keys.append( + ( + f"transformer.decoder.layers.{i}.cross_attn.out_proj.weight", + f"decoder.layers.{i}.encoder_attn.out_proj.weight", + ) + ) + rename_keys.append( + ( + f"transformer.decoder.layers.{i}.cross_attn.out_proj.bias", + f"decoder.layers.{i}.encoder_attn.out_proj.bias", + ) + ) + rename_keys.append((f"transformer.decoder.layers.{i}.linear1.weight", f"decoder.layers.{i}.fc1.weight")) + rename_keys.append((f"transformer.decoder.layers.{i}.linear1.bias", f"decoder.layers.{i}.fc1.bias")) + rename_keys.append((f"transformer.decoder.layers.{i}.linear2.weight", f"decoder.layers.{i}.fc2.weight")) + rename_keys.append((f"transformer.decoder.layers.{i}.linear2.bias", f"decoder.layers.{i}.fc2.bias")) + rename_keys.append( + (f"transformer.decoder.layers.{i}.norm1.weight", f"decoder.layers.{i}.self_attn_layer_norm.weight") + ) + rename_keys.append((f"transformer.decoder.layers.{i}.norm1.bias", f"decoder.layers.{i}.self_attn_layer_norm.bias")) + rename_keys.append( + (f"transformer.decoder.layers.{i}.norm2.weight", f"decoder.layers.{i}.encoder_attn_layer_norm.weight") + ) + rename_keys.append( + (f"transformer.decoder.layers.{i}.norm2.bias", f"decoder.layers.{i}.encoder_attn_layer_norm.bias") + ) + rename_keys.append((f"transformer.decoder.layers.{i}.norm3.weight", f"decoder.layers.{i}.final_layer_norm.weight")) + rename_keys.append((f"transformer.decoder.layers.{i}.norm3.bias", f"decoder.layers.{i}.final_layer_norm.bias")) + + # q, k, v projections in self/cross-attention in decoder for conditional DETR + # Self-attention projections moved into self_attn module + rename_keys.append( + ( + f"transformer.decoder.layers.{i}.sa_qcontent_proj.weight", + f"decoder.layers.{i}.self_attn.q_content_proj.weight", + ) + ) + rename_keys.append( + ( + f"transformer.decoder.layers.{i}.sa_kcontent_proj.weight", + f"decoder.layers.{i}.self_attn.k_content_proj.weight", + ) + ) + rename_keys.append( + (f"transformer.decoder.layers.{i}.sa_qpos_proj.weight", f"decoder.layers.{i}.self_attn.q_pos_proj.weight") + ) + rename_keys.append( + (f"transformer.decoder.layers.{i}.sa_kpos_proj.weight", f"decoder.layers.{i}.self_attn.k_pos_proj.weight") + ) + rename_keys.append( + (f"transformer.decoder.layers.{i}.sa_v_proj.weight", f"decoder.layers.{i}.self_attn.v_proj.weight") + ) + # Cross-attention projections moved into encoder_attn module + rename_keys.append( + ( + f"transformer.decoder.layers.{i}.ca_qcontent_proj.weight", + f"decoder.layers.{i}.encoder_attn.q_content_proj.weight", + ) + ) + # rename_keys.append((f"transformer.decoder.layers.{i}.ca_qpos_proj.weight", f"decoder.layers.{i}.encoder_attn.q_pos_proj.weight")) + rename_keys.append( + ( + f"transformer.decoder.layers.{i}.ca_kcontent_proj.weight", + f"decoder.layers.{i}.encoder_attn.k_content_proj.weight", + ) + ) + rename_keys.append( + (f"transformer.decoder.layers.{i}.ca_kpos_proj.weight", f"decoder.layers.{i}.encoder_attn.k_pos_proj.weight") + ) + rename_keys.append( + (f"transformer.decoder.layers.{i}.ca_v_proj.weight", f"decoder.layers.{i}.encoder_attn.v_proj.weight") + ) + rename_keys.append( + ( + f"transformer.decoder.layers.{i}.ca_qpos_sine_proj.weight", + f"decoder.layers.{i}.encoder_attn.q_pos_sine_proj.weight", + ) + ) + + rename_keys.append( + (f"transformer.decoder.layers.{i}.sa_qcontent_proj.bias", f"decoder.layers.{i}.self_attn.q_content_proj.bias") + ) + rename_keys.append( + (f"transformer.decoder.layers.{i}.sa_kcontent_proj.bias", f"decoder.layers.{i}.self_attn.k_content_proj.bias") + ) + rename_keys.append( + (f"transformer.decoder.layers.{i}.sa_qpos_proj.bias", f"decoder.layers.{i}.self_attn.q_pos_proj.bias") + ) + rename_keys.append( + (f"transformer.decoder.layers.{i}.sa_kpos_proj.bias", f"decoder.layers.{i}.self_attn.k_pos_proj.bias") + ) + rename_keys.append((f"transformer.decoder.layers.{i}.sa_v_proj.bias", f"decoder.layers.{i}.self_attn.v_proj.bias")) + rename_keys.append( + ( + f"transformer.decoder.layers.{i}.ca_qcontent_proj.bias", + f"decoder.layers.{i}.encoder_attn.q_content_proj.bias", + ) + ) + # rename_keys.append((f"transformer.decoder.layers.{i}.ca_qpos_proj.bias", f"decoder.layers.{i}.encoder_attn.q_pos_proj.bias")) + rename_keys.append( + ( + f"transformer.decoder.layers.{i}.ca_kcontent_proj.bias", + f"decoder.layers.{i}.encoder_attn.k_content_proj.bias", + ) + ) + rename_keys.append( + (f"transformer.decoder.layers.{i}.ca_kpos_proj.bias", f"decoder.layers.{i}.encoder_attn.k_pos_proj.bias") + ) + rename_keys.append( + (f"transformer.decoder.layers.{i}.ca_v_proj.bias", f"decoder.layers.{i}.encoder_attn.v_proj.bias") + ) + rename_keys.append( + ( + f"transformer.decoder.layers.{i}.ca_qpos_sine_proj.bias", + f"decoder.layers.{i}.encoder_attn.q_pos_sine_proj.bias", + ) + ) + +# convolutional projection + query embeddings + layernorm of decoder + class and bounding box heads +# for conditional DETR, also convert reference point head and query scale MLP +rename_keys.extend( + [ + ("input_proj.weight", "input_projection.weight"), + ("input_proj.bias", "input_projection.bias"), + ("query_embed.weight", "query_position_embeddings.weight"), + ("transformer.decoder.norm.weight", "decoder.layernorm.weight"), + ("transformer.decoder.norm.bias", "decoder.layernorm.bias"), + ("class_embed.weight", "class_labels_classifier.weight"), + ("class_embed.bias", "class_labels_classifier.bias"), + ("bbox_embed.layers.0.weight", "bbox_predictor.layers.0.weight"), + ("bbox_embed.layers.0.bias", "bbox_predictor.layers.0.bias"), + ("bbox_embed.layers.1.weight", "bbox_predictor.layers.1.weight"), + ("bbox_embed.layers.1.bias", "bbox_predictor.layers.1.bias"), + ("bbox_embed.layers.2.weight", "bbox_predictor.layers.2.weight"), + ("bbox_embed.layers.2.bias", "bbox_predictor.layers.2.bias"), + ("transformer.decoder.ref_point_head.layers.0.weight", "decoder.ref_point_head.layers.0.weight"), + ("transformer.decoder.ref_point_head.layers.0.bias", "decoder.ref_point_head.layers.0.bias"), + ("transformer.decoder.ref_point_head.layers.1.weight", "decoder.ref_point_head.layers.1.weight"), + ("transformer.decoder.ref_point_head.layers.1.bias", "decoder.ref_point_head.layers.1.bias"), + ("transformer.decoder.query_scale.layers.0.weight", "decoder.query_scale.layers.0.weight"), + ("transformer.decoder.query_scale.layers.0.bias", "decoder.query_scale.layers.0.bias"), + ("transformer.decoder.query_scale.layers.1.weight", "decoder.query_scale.layers.1.weight"), + ("transformer.decoder.query_scale.layers.1.bias", "decoder.query_scale.layers.1.bias"), + ("transformer.decoder.layers.0.ca_qpos_proj.weight", "decoder.layers.0.encoder_attn.q_pos_proj.weight"), + ("transformer.decoder.layers.0.ca_qpos_proj.bias", "decoder.layers.0.encoder_attn.q_pos_proj.bias"), + ] +) + + +def rename_key(state_dict, old, new): + val = state_dict.pop(old) + state_dict[new] = val + + +def rename_backbone_keys(state_dict): + new_state_dict = OrderedDict() + for key, value in state_dict.items(): + if "backbone.0.body" in key: + new_key = key.replace("backbone.0.body", "backbone.conv_encoder.model") + new_state_dict[new_key] = value + else: + new_state_dict[key] = value + + return new_state_dict + + +def read_in_q_k_v(state_dict, is_panoptic=False): + prefix = "" + if is_panoptic: + prefix = "conditional_detr." + + # first: transformer encoder + for i in range(6): + # read in weights + bias of input projection layer (in PyTorch's MultiHeadAttention, this is a single matrix + bias) + in_proj_weight = state_dict.pop(f"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_weight") + in_proj_bias = state_dict.pop(f"{prefix}transformer.encoder.layers.{i}.self_attn.in_proj_bias") + # next, add query, keys and values (in that order) to the state dict + state_dict[f"encoder.layers.{i}.self_attn.q_proj.weight"] = in_proj_weight[:256, :] + state_dict[f"encoder.layers.{i}.self_attn.q_proj.bias"] = in_proj_bias[:256] + state_dict[f"encoder.layers.{i}.self_attn.k_proj.weight"] = in_proj_weight[256:512, :] + state_dict[f"encoder.layers.{i}.self_attn.k_proj.bias"] = in_proj_bias[256:512] + state_dict[f"encoder.layers.{i}.self_attn.v_proj.weight"] = in_proj_weight[-256:, :] + state_dict[f"encoder.layers.{i}.self_attn.v_proj.bias"] = in_proj_bias[-256:] + + +# 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_conditional_detr_checkpoint(model_name, pytorch_dump_folder_path): + """ + Copy/paste/tweak model's weights to our CONDITIONAL_DETR structure. + """ + + # load default config + config = ConditionalDetrConfig() + # set backbone and dilation attributes + if "resnet101" in model_name: + config.backbone = "resnet101" + if "dc5" in model_name: + config.dilation = True + is_panoptic = "panoptic" in model_name + if is_panoptic: + config.num_labels = 250 + else: + config.num_labels = 91 + repo_id = "huggingface/label-files" + filename = "coco-detection-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 image processor + format = "coco_panoptic" if is_panoptic else "coco_detection" + image_processor = ConditionalDetrImageProcessor(format=format) + + # prepare image + img = prepare_img() + encoding = image_processor(images=img, return_tensors="pt") + pixel_values = encoding["pixel_values"] + + logger.info(f"Converting model {model_name}...") + + # load original model from torch hub + conditional_detr = torch.hub.load("DeppMeng/ConditionalDETR", model_name, pretrained=True).eval() + state_dict = conditional_detr.state_dict() + # rename keys + for src, dest in rename_keys: + if is_panoptic: + src = "conditional_detr." + src + rename_key(state_dict, src, dest) + state_dict = rename_backbone_keys(state_dict) + # query, key and value matrices need special treatment + read_in_q_k_v(state_dict, is_panoptic=is_panoptic) + # important: we need to prepend a prefix to each of the base model keys as the head models use different attributes for them + prefix = "conditional_detr.model." if is_panoptic else "model." + for key in state_dict.copy(): + if is_panoptic: + if ( + key.startswith("conditional_detr") + and not key.startswith("class_labels_classifier") + and not key.startswith("bbox_predictor") + ): + val = state_dict.pop(key) + state_dict["conditional_detr.model" + key[4:]] = val + elif "class_labels_classifier" in key or "bbox_predictor" in key: + val = state_dict.pop(key) + state_dict["conditional_detr." + key] = val + elif key.startswith("bbox_attention") or key.startswith("mask_head"): + continue + else: + val = state_dict.pop(key) + state_dict[prefix + key] = val + else: + if not key.startswith("class_labels_classifier") and not key.startswith("bbox_predictor"): + val = state_dict.pop(key) + state_dict[prefix + key] = val + # finally, create HuggingFace model and load state dict + model = ConditionalDetrForSegmentation(config) if is_panoptic else ConditionalDetrForObjectDetection(config) + model.load_state_dict(state_dict) + model.eval() + model.push_to_hub(repo_id=f"DepuMeng/{model_name}", commit_message="Add model") + # verify our conversion + original_outputs = conditional_detr(pixel_values) + outputs = model(pixel_values) + assert torch.allclose(outputs.logits, original_outputs["pred_logits"], atol=1e-4) + assert torch.allclose(outputs.pred_boxes, original_outputs["pred_boxes"], atol=1e-4) + if is_panoptic: + assert torch.allclose(outputs.pred_masks, original_outputs["pred_masks"], atol=1e-4) + + # Save model and image processor + logger.info(f"Saving PyTorch model and image processor to {pytorch_dump_folder_path}...") + Path(pytorch_dump_folder_path).mkdir(exist_ok=True) + model.save_pretrained(pytorch_dump_folder_path) + image_processor.save_pretrained(pytorch_dump_folder_path) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument( + "--model_name", + default="conditional_detr_resnet50", + type=str, + help="Name of the CONDITIONAL_DETR model you'd like to convert.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", default=None, type=str, help="Path to the folder to output PyTorch model." + ) + args = parser.parse_args() + convert_conditional_detr_checkpoint(args.model_name, args.pytorch_dump_folder_path) diff --git a/third_party/transformers/src/transformers/models/conditional_detr/image_processing_conditional_detr.py b/third_party/transformers/src/transformers/models/conditional_detr/image_processing_conditional_detr.py new file mode 100644 index 0000000000000000000000000000000000000000..539fe152f6061bcfbc719903a5cb81f0bd680976 --- /dev/null +++ b/third_party/transformers/src/transformers/models/conditional_detr/image_processing_conditional_detr.py @@ -0,0 +1,1083 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/conditional_detr/modular_conditional_detr.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_conditional_detr.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2022 Microsoft Research Asia 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 pathlib +from typing import Any, Optional + +import numpy as np +import torch +from torch import nn +from torchvision.io import read_image +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 ( + center_to_corners_format, + corners_to_center_format, + get_size_with_aspect_ratio, + safe_squeeze, +) +from ...image_utils import ( + IMAGENET_DEFAULT_MEAN, + IMAGENET_DEFAULT_STD, + AnnotationFormat, + AnnotationType, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, + get_image_size, + get_image_size_for_max_height_width, + get_max_height_width, + validate_annotations, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, logging + + +logger = logging.get_logger(__name__) + + +class ConditionalDetrImageProcessorKwargs(ImagesKwargs, total=False): + r""" + format (`str`, *optional*, defaults to `AnnotationFormat.COCO_DETECTION`): + Data format of the annotations. One of "coco_detection" or "coco_panoptic". + do_convert_annotations (`bool`, *optional*, defaults to `True`): + Controls whether to convert the annotations to the format expected by the CONDITIONAL_DETR model. Converts the + bounding boxes to the format `(center_x, center_y, width, height)` and in the range `[0, 1]`. + Can be overridden by the `do_convert_annotations` parameter in the `preprocess` method. + """ + + format: str | AnnotationFormat + do_convert_annotations: bool + + +SUPPORTED_ANNOTATION_FORMATS = (AnnotationFormat.COCO_DETECTION, AnnotationFormat.COCO_PANOPTIC) + + +def binary_mask_to_rle(mask): + """ + Converts given binary mask of shape `(height, width)` to the run-length encoding (RLE) format. + + Args: + mask (`torch.Tensor` or `numpy.array`): + A binary mask tensor of shape `(height, width)` where 0 denotes background and 1 denotes the target + segment_id or class_id. + Returns: + `List`: Run-length encoded list of the binary mask. Refer to COCO API for more information about the RLE + format. + """ + from ...utils import is_torch_tensor + + if is_torch_tensor(mask): + mask = mask.numpy() + + pixels = mask.flatten() + pixels = np.concatenate([[0], pixels, [0]]) + runs = np.where(pixels[1:] != pixels[:-1])[0] + 1 + runs[1::2] -= runs[::2] + return list(runs) + + +def convert_segmentation_to_rle(segmentation): + """ + Converts given segmentation map of shape `(height, width)` to the run-length encoding (RLE) format. + + Args: + segmentation (`torch.Tensor` or `numpy.array`): + A segmentation map of shape `(height, width)` where each value denotes a segment or class id. + Returns: + `list[List]`: A list of lists, where each list is the run-length encoding of a segment / class id. + """ + segment_ids = torch.unique(segmentation) + + run_length_encodings = [] + for idx in segment_ids: + mask = torch.where(segmentation == idx, 1, 0) + rle = binary_mask_to_rle(mask) + run_length_encodings.append(rle) + + return run_length_encodings + + +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_area = (mask_probs[k] >= mask_threshold).sum() + mask_exists = mask_k_area > 0 and original_area > 0 + + # Eliminate disconnected tiny segments + 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, mask_k + + +def compute_segments( + mask_probs, + pred_scores, + pred_labels, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + label_ids_to_fuse: set[int] | None = None, + 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.int32, device=mask_probs.device) + segments: list[dict] = [] + + if target_size is not None: + mask_probs = nn.functional.interpolate( + mask_probs.unsqueeze(0), size=target_size, mode="bilinear", align_corners=False + )[0] + + current_segment_id = 0 + + # Weigh each mask by its prediction score + mask_probs *= pred_scores.view(-1, 1, 1) + mask_labels = mask_probs.argmax(0) # [height, width] + + # Keep track of instances of each class + stuff_memory_list: dict[str, int] = {} + for k in range(pred_labels.shape[0]): + pred_class = pred_labels[k].item() + should_fuse = pred_class in label_ids_to_fuse + + # Check if mask exists and large enough to be a segment + mask_exists, mask_k = check_segment_validity( + mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold + ) + + if mask_exists: + if pred_class in stuff_memory_list: + current_segment_id = stuff_memory_list[pred_class] + else: + current_segment_id += 1 + + # Add current object segment to final segmentation map + segmentation[mask_k] = current_segment_id + segment_score = round(pred_scores[k].item(), 6) + segments.append( + { + "id": current_segment_id, + "label_id": pred_class, + "was_fused": should_fuse, + "score": segment_score, + } + ) + if should_fuse: + stuff_memory_list[pred_class] = current_segment_id + + return segmentation, segments + + +# inspired by https://github.com/facebookresearch/conditional_detr/blob/master/datasets/coco.py#L33 +def convert_coco_poly_to_mask(segmentations, height: int, width: int, device: torch.device) -> torch.Tensor: + """ + Convert a COCO polygon annotation to a mask. + + Args: + segmentations (`list[list[float]]`): + List of polygons, each polygon represented by a list of x-y coordinates. + height (`int`): + Height of the mask. + width (`int`): + Width of the mask. + """ + try: + from pycocotools import mask as coco_mask + except ImportError: + raise ImportError("Pycocotools is not installed in your environment.") + + masks = [] + for polygons in segmentations: + rles = coco_mask.frPyObjects(polygons, height, width) + mask = coco_mask.decode(rles) + if len(mask.shape) < 3: + mask = mask[..., None] + mask = torch.as_tensor(mask, dtype=torch.uint8, device=device) + mask = torch.any(mask, axis=2) + masks.append(mask) + if masks: + masks = torch.stack(masks, axis=0) + else: + masks = torch.zeros((0, height, width), dtype=torch.uint8, device=device) + + return masks + + +# inspired by https://github.com/facebookresearch/conditional_detr/blob/master/datasets/coco.py#L50 +def prepare_coco_detection_annotation( + image, + target, + return_segmentation_masks: bool = False, + input_data_format: ChannelDimension | str | None = None, +): + """ + Convert the target in COCO format into the format expected by CONDITIONAL_DETR. + """ + image_height, image_width = image.size()[-2:] + + image_id = target["image_id"] + image_id = torch.as_tensor([image_id], dtype=torch.int64, device=image.device) + + # Get all COCO annotations for the given image. + annotations = target["annotations"] + classes = [] + area = [] + boxes = [] + keypoints = [] + for obj in annotations: + if "iscrowd" not in obj or obj["iscrowd"] == 0: + classes.append(obj["category_id"]) + area.append(obj["area"]) + boxes.append(obj["bbox"]) + if "keypoints" in obj: + keypoints.append(obj["keypoints"]) + + classes = torch.as_tensor(classes, dtype=torch.int64, device=image.device) + area = torch.as_tensor(area, dtype=torch.float32, device=image.device) + iscrowd = torch.zeros_like(classes, dtype=torch.int64, device=image.device) + # guard against no boxes via resizing + boxes = torch.as_tensor(boxes, dtype=torch.float32, device=image.device).reshape(-1, 4) + boxes[:, 2:] += boxes[:, :2] + boxes[:, 0::2] = boxes[:, 0::2].clip(min=0, max=image_width) + boxes[:, 1::2] = boxes[:, 1::2].clip(min=0, max=image_height) + + keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0]) + + new_target = { + "image_id": image_id, + "class_labels": classes[keep], + "boxes": boxes[keep], + "area": area[keep], + "iscrowd": iscrowd[keep], + "orig_size": torch.as_tensor([int(image_height), int(image_width)], dtype=torch.int64, device=image.device), + } + + if keypoints: + keypoints = torch.as_tensor(keypoints, dtype=torch.float32, device=image.device) + # Apply the keep mask here to filter the relevant annotations + keypoints = keypoints[keep] + num_keypoints = keypoints.shape[0] + keypoints = keypoints.reshape((-1, 3)) if num_keypoints else keypoints + new_target["keypoints"] = keypoints + + if return_segmentation_masks: + segmentation_masks = [obj["segmentation"] for obj in annotations] + masks = convert_coco_poly_to_mask(segmentation_masks, image_height, image_width, device=image.device) + new_target["masks"] = masks[keep] + + return new_target + + +def masks_to_boxes(masks: torch.Tensor) -> torch.Tensor: + """ + Compute the bounding boxes around the provided panoptic segmentation masks. + + Args: + masks: masks in format `[number_masks, height, width]` where N is the number of masks + + Returns: + boxes: bounding boxes in format `[number_masks, 4]` in xyxy format + """ + if masks.numel() == 0: + return torch.zeros((0, 4), device=masks.device) + + h, w = masks.shape[-2:] + y = torch.arange(0, h, dtype=torch.float32, device=masks.device) + x = torch.arange(0, w, dtype=torch.float32, device=masks.device) + # see https://github.com/pytorch/pytorch/issues/50276 + y, x = torch.meshgrid(y, x, indexing="ij") + + x_mask = masks * torch.unsqueeze(x, 0) + x_max = x_mask.view(x_mask.shape[0], -1).max(-1)[0] + x_min = ( + torch.where(masks, x.unsqueeze(0), torch.tensor(1e8, device=masks.device)).view(masks.shape[0], -1).min(-1)[0] + ) + + y_mask = masks * torch.unsqueeze(y, 0) + y_max = y_mask.view(y_mask.shape[0], -1).max(-1)[0] + y_min = ( + torch.where(masks, y.unsqueeze(0), torch.tensor(1e8, device=masks.device)).view(masks.shape[0], -1).min(-1)[0] + ) + + return torch.stack([x_min, y_min, x_max, y_max], 1) + + +# 2 functions below adapted from https://github.com/cocodataset/panopticapi/blob/master/panopticapi/utils.py +# Copyright (c) 2018, Alexander Kirillov +# All rights reserved. +def rgb_to_id(color): + """ + Converts RGB color to unique ID. + """ + if isinstance(color, torch.Tensor) and len(color.shape) == 3: + if color.dtype == torch.uint8: + color = color.to(torch.int32) + return color[:, :, 0] + 256 * color[:, :, 1] + 256 * 256 * color[:, :, 2] + return int(color[0] + 256 * color[1] + 256 * 256 * color[2]) + + +def prepare_coco_panoptic_annotation( + image: torch.Tensor, + target: dict, + masks_path: str | pathlib.Path, + return_masks: bool = True, + input_data_format: ChannelDimension | str = None, +) -> dict: + """ + Prepare a coco panoptic annotation for CONDITIONAL_DETR. + """ + image_height, image_width = get_image_size(image, channel_dim=input_data_format) + annotation_path = pathlib.Path(masks_path) / target["file_name"] + + new_target = {} + new_target["image_id"] = torch.as_tensor( + [target["image_id"] if "image_id" in target else target["id"]], dtype=torch.int64, device=image.device + ) + new_target["size"] = torch.as_tensor([image_height, image_width], dtype=torch.int64, device=image.device) + new_target["orig_size"] = torch.as_tensor([image_height, image_width], dtype=torch.int64, device=image.device) + + if "segments_info" in target: + masks = read_image(annotation_path).permute(1, 2, 0).to(dtype=torch.int32, device=image.device) + masks = rgb_to_id(masks) + + ids = torch.as_tensor([segment_info["id"] for segment_info in target["segments_info"]], device=image.device) + masks = masks == ids[:, None, None] + masks = masks.to(torch.bool) + if return_masks: + new_target["masks"] = masks + new_target["boxes"] = masks_to_boxes(masks) + new_target["class_labels"] = torch.as_tensor( + [segment_info["category_id"] for segment_info in target["segments_info"]], + dtype=torch.int64, + device=image.device, + ) + new_target["iscrowd"] = torch.as_tensor( + [segment_info["iscrowd"] for segment_info in target["segments_info"]], + dtype=torch.int64, + device=image.device, + ) + new_target["area"] = torch.as_tensor( + [segment_info["area"] for segment_info in target["segments_info"]], + dtype=torch.float32, + device=image.device, + ) + + return new_target + + +@auto_docstring +class ConditionalDetrImageProcessor(TorchvisionBackend): + valid_kwargs = ConditionalDetrImageProcessorKwargs + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_DEFAULT_MEAN + image_std = IMAGENET_DEFAULT_STD + format = AnnotationFormat.COCO_DETECTION + do_resize = True + do_rescale = True + do_normalize = True + do_pad = True + size = {"shortest_edge": 800, "longest_edge": 1333} + default_to_square = False + model_input_names = ["pixel_values", "pixel_mask"] + + def __init__(self, **kwargs: Unpack[ConditionalDetrImageProcessorKwargs]) -> None: + kwargs.setdefault("do_pad", kwargs.pop("pad_and_return_pixel_mask", self.do_pad)) + + size = kwargs.pop("size", None) + max_size = None if size is None else kwargs.pop("max_size", 1333) + size = size if size is not None else {"shortest_edge": 800, "longest_edge": 1333} + # Convert size dict for backwards compat with max_size parameter + kwargs["size"] = get_size_dict(size, max_size=max_size, default_to_square=False) + + # Backwards compatibility + do_convert_annotations = kwargs.get("do_convert_annotations") + do_normalize = kwargs.get("do_normalize") + if do_convert_annotations is None and getattr(self, "do_convert_annotations", None) is None: + self.do_convert_annotations = do_normalize if do_normalize is not None else self.do_normalize + + super().__init__(**kwargs) + + def prepare_annotation( + self, + image: torch.Tensor, + target: dict, + format: AnnotationFormat | None = None, + return_segmentation_masks: bool | None = None, + masks_path: str | pathlib.Path | None = None, + input_data_format: str | ChannelDimension | None = None, + ) -> dict: + """ + Prepare an annotation for feeding into CONDITIONAL_DETR model. + """ + format = format if format is not None else self.format + + if format == AnnotationFormat.COCO_DETECTION: + return_segmentation_masks = False if return_segmentation_masks is None else return_segmentation_masks + target = prepare_coco_detection_annotation( + image, target, return_segmentation_masks, input_data_format=input_data_format + ) + elif format == AnnotationFormat.COCO_PANOPTIC: + return_segmentation_masks = True if return_segmentation_masks is None else return_segmentation_masks + target = prepare_coco_panoptic_annotation( + image, + target, + masks_path=masks_path, + return_masks=return_segmentation_masks, + input_data_format=input_data_format, + ) + else: + raise ValueError(f"Format {format} is not supported.") + return target + + def resize( + self, + image: torch.Tensor, + size: SizeDict, + resample: Optional["PILImageResampling | tvF.InterpolationMode | int"] = None, + **kwargs, + ) -> torch.Tensor: + """ + Resize the image to the given size. Size can be `min_size` (scalar) or `(height, width)` tuple. If size is an + int, smaller edge of the image will be matched to this number. + + Args: + image (`torch.Tensor`): + Image to resize. + size (`SizeDict`): + Size of the image's `(height, width)` dimensions after resizing. Available options are: + - `{"height": int, "width": int}`: The image will be resized to the exact size `(height, width)`. + Do NOT keep the aspect ratio. + - `{"shortest_edge": int, "longest_edge": int}`: The image will be resized to a maximum size respecting + the aspect ratio and keeping the shortest edge less or equal to `shortest_edge` and the longest edge + less or equal to `longest_edge`. + - `{"max_height": int, "max_width": int}`: The image will be resized to the maximum size respecting the + aspect ratio and keeping the height less or equal to `max_height` and the width less or equal to + `max_width`. + resample (`PILImageResampling | tvF.InterpolationMode | int`, *optional*, defaults to `PILImageResampling.BILINEAR`): + Resampling filter to use if resizing the image. + """ + if size.shortest_edge and size.longest_edge: + # Resize the image so that the shortest edge or the longest edge is of the given size + # while maintaining the aspect ratio of the original image. + new_size = get_size_with_aspect_ratio(image.shape[-2:], size.shortest_edge, size.longest_edge) + elif size.max_height and size.max_width: + new_size = get_image_size_for_max_height_width(image.shape[-2:], size.max_height, size.max_width) + elif size.height and size.width: + new_size = (size.height, size.width) + else: + raise ValueError( + f"Size must contain 'height' and 'width' keys or 'shortest_edge' and 'longest_edge' keys. Got {size}." + ) + + image = super().resize( + image, size=SizeDict(height=new_size[0], width=new_size[1]), resample=resample, **kwargs + ) + return image + + def resize_annotation( + self, + annotation: dict[str, Any], + orig_size: tuple[int, int], + target_size: tuple[int, int], + threshold: float = 0.5, + resample: Optional["PILImageResampling | tvF.InterpolationMode | int"] = PILImageResampling.NEAREST, + ): + """ + Resizes an annotation to a target size. + + Args: + annotation (`dict[str, Any]`): + The annotation dictionary. + orig_size (`tuple[int, int]`): + The original size of the input image. + target_size (`tuple[int, int]`): + The target size of the image, as returned by the preprocessing `resize` step. + threshold (`float`, *optional*, defaults to 0.5): + The threshold used to binarize the segmentation masks. + resample (`PILImageResampling | tvF.InterpolationMode | int`, defaults to `tvF.InterpolationMode.NEAREST_EXACT`): + The resampling filter to use when resizing the masks. + """ + ratio_height, ratio_width = [target / orig for target, orig in zip(target_size, orig_size)] + + new_annotation = {} + new_annotation["size"] = target_size + + for key, value in annotation.items(): + if key == "boxes": + boxes = value + scaled_boxes = boxes * torch.as_tensor( + [ratio_width, ratio_height, ratio_width, ratio_height], dtype=torch.float32, device=boxes.device + ) + new_annotation["boxes"] = scaled_boxes + elif key == "area": + area = value + scaled_area = area * (ratio_width * ratio_height) + new_annotation["area"] = scaled_area + elif key == "masks": + masks = value[:, None] + masks = [ + super(ConditionalDetrImageProcessor, self).resize( + mask, size=SizeDict(height=target_size[0], width=target_size[1]), resample=resample + ) + for mask in masks + ] + masks = torch.stack(masks).to(torch.float32) + masks = masks[:, 0] > threshold + new_annotation["masks"] = masks + elif key == "size": + new_annotation["size"] = target_size + else: + new_annotation[key] = value + + return new_annotation + + def normalize_annotation(self, annotation: dict, image_size: tuple[int, int]) -> dict: + image_height, image_width = image_size + norm_annotation = {} + for key, value in annotation.items(): + if key == "boxes": + boxes = value + boxes = corners_to_center_format(boxes) + boxes /= torch.as_tensor( + [image_width, image_height, image_width, image_height], dtype=torch.float32, device=boxes.device + ) + norm_annotation[key] = boxes + else: + norm_annotation[key] = value + return norm_annotation + + def _update_annotation_for_padded_image( + self, + annotation: dict, + input_image_size: tuple[int, int], + output_image_size: tuple[int, int], + padding, + update_bboxes, + ) -> dict: + """ + Update the annotation for a padded image. + """ + new_annotation = {} + new_annotation["size"] = output_image_size + ratio_height, ratio_width = (input / output for output, input in zip(output_image_size, input_image_size)) + + for key, value in annotation.items(): + if key == "masks": + masks = value + masks = tvF.pad( + masks, + padding, + fill=0, + ) + masks = safe_squeeze(masks, 1) + new_annotation["masks"] = masks + elif key == "boxes" and update_bboxes: + boxes = value + boxes *= torch.as_tensor([ratio_width, ratio_height, ratio_width, ratio_height], device=boxes.device) + new_annotation["boxes"] = boxes + elif key == "size": + new_annotation["size"] = output_image_size + else: + new_annotation[key] = value + return new_annotation + + def pad( + self, + image: torch.Tensor, + padded_size: tuple[int, int], + annotation: dict[str, Any] | None = None, + update_bboxes: bool = True, + fill: int = 0, + ): + original_size = image.size()[-2:] + padding_bottom = padded_size[0] - original_size[0] + padding_right = padded_size[1] - original_size[1] + if padding_bottom < 0 or padding_right < 0: + raise ValueError( + f"Padding dimensions are negative. Please make sure that the padded size is larger than the " + f"original size. Got padded size: {padded_size}, original size: {original_size}." + ) + if original_size != padded_size: + padding = [0, 0, padding_right, padding_bottom] + image = tvF.pad(image, padding, fill=fill) + if annotation is not None: + annotation = self._update_annotation_for_padded_image( + annotation, original_size, padded_size, padding, update_bboxes + ) + + # Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding. + pixel_mask = torch.zeros(padded_size, dtype=torch.int64, device=image.device) + pixel_mask[: original_size[0], : original_size[1]] = 1 + + return image, pixel_mask, annotation + + @auto_docstring + def preprocess( + self, + images: ImageInput, + annotations: AnnotationType | list[AnnotationType] | None = None, + return_segmentation_masks: bool | None = None, + masks_path: str | pathlib.Path | None = None, + **kwargs: Unpack[ConditionalDetrImageProcessorKwargs], + ) -> BatchFeature: + r""" + annotations (`AnnotationType` or `list[AnnotationType]`, *optional*): + Annotations to transform according to the padding that is applied to the images. + return_segmentation_masks (`bool`, *optional*, defaults to `self.return_segmentation_masks`): + Whether to return segmentation masks. + masks_path (`str` or `pathlib.Path`, *optional*): + Path to the directory containing the segmentation masks. + """ + return super().preprocess(images, annotations, return_segmentation_masks, masks_path, **kwargs) + + def _preprocess( + self, + images: list["torch.Tensor"], + annotations: AnnotationType | list[AnnotationType] | None, + return_segmentation_masks: bool, + masks_path: str | pathlib.Path | None, + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + do_convert_annotations: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + do_pad: bool, + pad_size: SizeDict | None, + format: str | AnnotationFormat | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + """ + Preprocess an image or a batch of images so that it can be used by the model. + """ + if annotations is not None and isinstance(annotations, dict): + annotations = [annotations] + + if annotations is not None and len(images) != len(annotations): + raise ValueError( + f"The number of images ({len(images)}) and annotations ({len(annotations)}) do not match." + ) + + format = AnnotationFormat(format) + if annotations is not None: + validate_annotations(format, SUPPORTED_ANNOTATION_FORMATS, annotations) + + if ( + masks_path is not None + and format == AnnotationFormat.COCO_PANOPTIC + and not isinstance(masks_path, (pathlib.Path, str)) + ): + raise ValueError( + "The path to the directory containing the mask PNG files should be provided as a" + f" `pathlib.Path` or string object, but is {type(masks_path)} instead." + ) + + data = {} + + processed_images = [] + processed_annotations = [] + pixel_masks = [] # Initialize pixel_masks here + for image, annotation in zip(images, annotations if annotations is not None else [None] * len(images)): + # prepare (COCO annotations as a list of Dict -> CONDITIONAL_DETR target as a single Dict per image) + if annotations is not None: + annotation = self.prepare_annotation( + image, + annotation, + format, + return_segmentation_masks=return_segmentation_masks, + masks_path=masks_path, + input_data_format=ChannelDimension.FIRST, + ) + + if do_resize: + resized_image = self.resize(image, size=size, resample=resample) + if annotations is not None: + annotation = self.resize_annotation( + annotation, + orig_size=image.size()[-2:], + target_size=resized_image.size()[-2:], + ) + image = resized_image + # Fused rescale and normalize + image = self.rescale_and_normalize(image, do_rescale, rescale_factor, do_normalize, image_mean, image_std) + if do_convert_annotations and annotations is not None: + annotation = self.normalize_annotation(annotation, get_image_size(image, ChannelDimension.FIRST)) + + processed_images.append(image) + processed_annotations.append(annotation) + images = processed_images + annotations = processed_annotations if annotations is not None else None + + if do_pad: + # depends on all resized image shapes so we need another loop + if pad_size is not None: + padded_size = (pad_size.height, pad_size.width) + else: + padded_size = get_max_height_width(images) + + padded_images = [] + padded_annotations = [] + for image, annotation in zip(images, annotations if annotations is not None else [None] * len(images)): + # Pads images and returns their mask: {'pixel_values': ..., 'pixel_mask': ...} + if padded_size == image.size()[-2:]: + padded_images.append(image) + pixel_masks.append(torch.ones(padded_size, dtype=torch.int64, device=image.device)) + padded_annotations.append(annotation) + continue + image, pixel_mask, annotation = self.pad( + image, padded_size, annotation=annotation, update_bboxes=do_convert_annotations + ) + padded_images.append(image) + padded_annotations.append(annotation) + pixel_masks.append(pixel_mask) + images = padded_images + annotations = padded_annotations if annotations is not None else None + data.update({"pixel_mask": torch.stack(pixel_masks, dim=0)}) + + data.update({"pixel_values": torch.stack(images, dim=0)}) + encoded_inputs = BatchFeature(data, tensor_type=return_tensors) + if annotations is not None: + encoded_inputs["labels"] = [ + BatchFeature(annotation, tensor_type=return_tensors) for annotation in annotations + ] + return encoded_inputs + + def post_process_object_detection( + self, outputs, threshold: float = 0.5, target_sizes: TensorType | list[tuple] = None, top_k: int = 100 + ): + """ + Converts the raw output of [`ConditionalDetrForObjectDetection`] into final bounding boxes in (top_left_x, + top_left_y, bottom_right_x, bottom_right_y) format. Only supports PyTorch. + + Args: + outputs ([`ConditionalDetrObjectDetectionOutput`]): + Raw outputs of the model. + threshold (`float`, *optional*): + Score threshold to keep object detection predictions. + target_sizes (`torch.Tensor` or `list[tuple[int, int]]`, *optional*): + Tensor of shape `(batch_size, 2)` or list of tuples (`tuple[int, int]`) containing the target size + (height, width) of each image in the batch. If left to None, predictions will not be resized. + top_k (`int`, *optional*, defaults to 100): + Keep only top k bounding boxes before filtering by thresholding. + + Returns: + `list[Dict]`: A list of dictionaries, each dictionary containing the scores, labels and boxes for an image + in the batch as predicted by the model. + """ + out_logits, out_bbox = outputs.logits, outputs.pred_boxes + + if target_sizes is not None: + if len(out_logits) != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + prob = out_logits.sigmoid() + prob = prob.view(out_logits.shape[0], -1) + k_value = min(top_k, prob.size(1)) + topk_values, topk_indexes = torch.topk(prob, k_value, dim=1) + scores = topk_values + topk_boxes = torch.div(topk_indexes, out_logits.shape[2], rounding_mode="floor") + labels = topk_indexes % out_logits.shape[2] + boxes = center_to_corners_format(out_bbox) + boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4)) + + # and from relative [0, 1] to absolute [0, height] coordinates + if target_sizes is not None: + if isinstance(target_sizes, list): + img_h = torch.Tensor([i[0] for i in target_sizes]) + img_w = torch.Tensor([i[1] for i in target_sizes]) + else: + img_h, img_w = target_sizes.unbind(1) + scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(boxes.device) + boxes = boxes * scale_fct[:, None, :] + + results = [] + for s, l, b in zip(scores, labels, boxes): + score = s[s > threshold] + label = l[s > threshold] + box = b[s > threshold] + results.append({"scores": score, "labels": label, "boxes": box}) + + return results + + def post_process_semantic_segmentation(self, outputs, target_sizes: list[tuple[int, int]] | None = None): + """ + Converts the output of [`ConditionalDetrForSegmentation`] into semantic segmentation maps. Only supports PyTorch. + + Args: + outputs ([`ConditionalDetrForSegmentation`]): + Raw outputs of the model. + target_sizes (`list[tuple[int, int]]`, *optional*): + A list of tuples (`tuple[int, int]`) containing the target size (height, width) of each image in the + batch. If unset, predictions will not be resized. + Returns: + `list[torch.Tensor]`: + A list of length `batch_size`, where each item is a semantic segmentation map of shape (height, width) + corresponding to the target_sizes entry (if `target_sizes` is specified). Each entry of each + `torch.Tensor` correspond to a semantic class id. + """ + class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes] + masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width] + + # Conditional DETR does not have a null class, so we use all classes + masks_classes = class_queries_logits.softmax(dim=-1) + masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # Semantic segmentation logits of shape (batch_size, num_classes, height, width) + segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs) + batch_size = class_queries_logits.shape[0] + + # Resize logits and compute semantic segmentation maps + if target_sizes is not None: + if batch_size != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + semantic_segmentation = [] + for idx in range(batch_size): + resized_logits = nn.functional.interpolate( + segmentation[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False + ) + semantic_map = resized_logits[0].argmax(dim=0) + semantic_segmentation.append(semantic_map) + else: + semantic_segmentation = segmentation.argmax(dim=1) + semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])] + + return semantic_segmentation + + def post_process_instance_segmentation( + self, + outputs, + threshold: float = 0.5, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + target_sizes: list[tuple[int, int]] | None = None, + return_coco_annotation: bool | None = False, + ) -> list[dict]: + """ + Converts the output of [`ConditionalDetrForSegmentation`] into instance segmentation predictions. Only supports PyTorch. + + Args: + outputs ([`ConditionalDetrForSegmentation`]): + Raw outputs of the model. + threshold (`float`, *optional*, defaults to 0.5): + The probability score threshold to keep predicted instance masks. + mask_threshold (`float`, *optional*, defaults to 0.5): + Threshold to use when turning the predicted masks into binary values. + overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8): + The overlap mask area threshold to merge or discard small disconnected parts within each binary + instance mask. + target_sizes (`list[Tuple]`, *optional*): + List of length (batch_size), where each list item (`tuple[int, int]]`) corresponds to the requested + final size (height, width) of each prediction. If unset, predictions will not be resized. + return_coco_annotation (`bool`, *optional*): + Defaults to `False`. If set to `True`, segmentation maps are returned in COCO run-length encoding (RLE) + format. + Returns: + `list[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys: + - **segmentation** -- A tensor of shape `(height, width)` where each pixel represents a `segment_id` or + `list[List]` run-length encoding (RLE) of the segmentation map if return_coco_annotation is set to + `True`. Set to `None` if no mask if found above `threshold`. + - **segments_info** -- A dictionary that contains additional information on each segment. + - **id** -- An integer representing the `segment_id`. + - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`. + - **score** -- Prediction score of segment with `segment_id`. + """ + class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width] + + batch_size = class_queries_logits.shape[0] + num_labels = class_queries_logits.shape[-1] - 1 + + mask_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # Predicted label and score of each query (batch_size, num_queries) + pred_scores, pred_labels = nn.functional.softmax(class_queries_logits, dim=-1).max(-1) + + # Loop over items in batch size + results: list[dict[str, TensorType]] = [] + + for i in range(batch_size): + mask_probs_item, pred_scores_item, pred_labels_item = remove_low_and_no_objects( + mask_probs[i], pred_scores[i], pred_labels[i], threshold, num_labels + ) + + # No mask found + if mask_probs_item.shape[0] <= 0: + height, width = target_sizes[i] if target_sizes is not None else mask_probs_item.shape[1:] + segmentation = torch.zeros((height, width)) - 1 + results.append({"segmentation": segmentation, "segments_info": []}) + continue + + # Get segmentation map and segment information of batch item + target_size = target_sizes[i] if target_sizes is not None else None + segmentation, segments = compute_segments( + mask_probs=mask_probs_item, + pred_scores=pred_scores_item, + pred_labels=pred_labels_item, + mask_threshold=mask_threshold, + overlap_mask_area_threshold=overlap_mask_area_threshold, + label_ids_to_fuse=[], + target_size=target_size, + ) + + # Return segmentation map in run-length encoding (RLE) format + if return_coco_annotation: + segmentation = convert_segmentation_to_rle(segmentation) + + results.append({"segmentation": segmentation, "segments_info": segments}) + return results + + def post_process_panoptic_segmentation( + self, + outputs, + threshold: float = 0.5, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + label_ids_to_fuse: set[int] | None = None, + target_sizes: list[tuple[int, int]] | None = None, + ) -> list[dict]: + """ + Converts the output of [`ConditionalDetrForSegmentation`] into image panoptic segmentation predictions. Only supports + PyTorch. + + Args: + outputs ([`ConditionalDetrForSegmentation`]): + The outputs from [`ConditionalDetrForSegmentation`]. + threshold (`float`, *optional*, defaults to 0.5): + The probability score threshold to keep predicted instance masks. + mask_threshold (`float`, *optional*, defaults to 0.5): + Threshold to use when turning the predicted masks into binary values. + overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8): + The overlap mask area threshold to merge or discard small disconnected parts within each binary + instance mask. + label_ids_to_fuse (`Set[int]`, *optional*): + The labels in this state will have all their instances be fused together. For instance we could say + there can only be one sky in an image, but several persons, so the label ID for sky would be in that + set, but not the one for person. + target_sizes (`list[Tuple]`, *optional*): + List of length (batch_size), where each list item (`tuple[int, int]]`) corresponds to the requested + final size (height, width) of each prediction in batch. If unset, predictions will not be resized. + Returns: + `list[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys: + - **segmentation** -- a tensor of shape `(height, width)` where each pixel represents a `segment_id` or + `None` if no mask if found above `threshold`. If `target_sizes` is specified, segmentation is resized to + the corresponding `target_sizes` entry. + - **segments_info** -- A dictionary that contains additional information on each segment. + - **id** -- an integer representing the `segment_id`. + - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`. + - **was_fused** -- a boolean, `True` if `label_id` was in `label_ids_to_fuse`, `False` otherwise. + Multiple instances of the same class / label were fused and assigned a single `segment_id`. + - **score** -- Prediction score of segment with `segment_id`. + """ + + if label_ids_to_fuse is None: + logger.warning_once("`label_ids_to_fuse` unset. No instance will be fused.") + label_ids_to_fuse = set() + + class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width] + + batch_size = class_queries_logits.shape[0] + num_labels = class_queries_logits.shape[-1] - 1 + + mask_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # Predicted label and score of each query (batch_size, num_queries) + pred_scores, pred_labels = nn.functional.softmax(class_queries_logits, dim=-1).max(-1) + + # Loop over items in batch size + results: list[dict[str, TensorType]] = [] + + for i in range(batch_size): + mask_probs_item, pred_scores_item, pred_labels_item = remove_low_and_no_objects( + mask_probs[i], pred_scores[i], pred_labels[i], threshold, num_labels + ) + + # No mask found + if mask_probs_item.shape[0] <= 0: + height, width = target_sizes[i] if target_sizes is not None else mask_probs_item.shape[1:] + segmentation = torch.zeros((height, width)) - 1 + results.append({"segmentation": segmentation, "segments_info": []}) + continue + + # Get segmentation map and segment information of batch item + target_size = target_sizes[i] if target_sizes is not None else None + segmentation, segments = compute_segments( + mask_probs=mask_probs_item, + pred_scores=pred_scores_item, + pred_labels=pred_labels_item, + mask_threshold=mask_threshold, + overlap_mask_area_threshold=overlap_mask_area_threshold, + label_ids_to_fuse=label_ids_to_fuse, + target_size=target_size, + ) + + results.append({"segmentation": segmentation, "segments_info": segments}) + return results + + +__all__ = ["ConditionalDetrImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/conditional_detr/image_processing_pil_conditional_detr.py b/third_party/transformers/src/transformers/models/conditional_detr/image_processing_pil_conditional_detr.py new file mode 100644 index 0000000000000000000000000000000000000000..359c4c706f7c6a9f34ab42df22c92ec0bf73ac95 --- /dev/null +++ b/third_party/transformers/src/transformers/models/conditional_detr/image_processing_pil_conditional_detr.py @@ -0,0 +1,1136 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/conditional_detr/modular_conditional_detr.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_conditional_detr.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2022 Microsoft Research Asia 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 pathlib +from typing import Any, Optional + +import numpy as np + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import ( + PaddingMode, + center_to_corners_format, + corners_to_center_format, + get_size_with_aspect_ratio, + pad, + resize, + safe_squeeze, +) +from ...image_utils import ( + IMAGENET_DEFAULT_MEAN, + IMAGENET_DEFAULT_STD, + AnnotationFormat, + AnnotationType, + ChannelDimension, + ImageInput, + PILImageResampling, + SizeDict, + get_image_size, + get_image_size_for_max_height_width, + get_max_height_width, + validate_annotations, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring, is_torch_available, is_vision_available, logging, requires_backends +from ...utils.import_utils import requires + + +if is_vision_available(): + import PIL.Image +if is_torch_available(): + import torch + from torch import nn + +logger = logging.get_logger(__name__) + + +class ConditionalDetrImageProcessorKwargs(ImagesKwargs, total=False): + r""" + format (`str`, *optional*, defaults to `AnnotationFormat.COCO_DETECTION`): + Data format of the annotations. One of "coco_detection" or "coco_panoptic". + do_convert_annotations (`bool`, *optional*, defaults to `True`): + Controls whether to convert the annotations to the format expected by the CONDITIONAL_DETR model. Converts the + bounding boxes to the format `(center_x, center_y, width, height)` and in the range `[0, 1]`. + Can be overridden by the `do_convert_annotations` parameter in the `preprocess` method. + """ + + format: str | AnnotationFormat + do_convert_annotations: bool + + +SUPPORTED_ANNOTATION_FORMATS = (AnnotationFormat.COCO_DETECTION, AnnotationFormat.COCO_PANOPTIC) + + +# inspired by https://github.com/facebookresearch/conditional_detr/blob/master/datasets/coco.py#L33 +def convert_coco_poly_to_mask(segmentations, height: int, width: int) -> np.ndarray: + """ + Convert a COCO polygon annotation to a mask. + + Args: + segmentations (`list[list[float]]`): + List of polygons, each polygon represented by a list of x-y coordinates. + height (`int`): + Height of the mask. + width (`int`): + Width of the mask. + """ + try: + from pycocotools import mask as coco_mask + except ImportError: + raise ImportError("Pycocotools is not installed in your environment.") + + masks = [] + for polygons in segmentations: + rles = coco_mask.frPyObjects(polygons, height, width) + mask = coco_mask.decode(rles) + if len(mask.shape) < 3: + mask = mask[..., None] + mask = np.asarray(mask, dtype=np.uint8) + mask = np.any(mask, axis=2) + masks.append(mask) + if masks: + masks = np.stack(masks, axis=0) + else: + masks = np.zeros((0, height, width), dtype=np.uint8) + + return masks + + +# inspired by https://github.com/facebookresearch/conditional_detr/blob/master/datasets/coco.py#L50 +def prepare_coco_detection_annotation( + image, + target, + return_segmentation_masks: bool = False, + input_data_format: ChannelDimension | str | None = None, +): + """ + Convert the target in COCO format into the format expected by CONDITIONAL_DETR. + """ + image_height, image_width = get_image_size(image, channel_dim=input_data_format) + + image_id = target["image_id"] + image_id = np.asarray([image_id], dtype=np.int64) + + # Get all COCO annotations for the given image. + annotations = target["annotations"] + annotations = [obj for obj in annotations if "iscrowd" not in obj or obj["iscrowd"] == 0] + + classes = [obj["category_id"] for obj in annotations] + classes = np.asarray(classes, dtype=np.int64) + + # for conversion to coco api + area = np.asarray([obj["area"] for obj in annotations], dtype=np.float32) + iscrowd = np.asarray([obj.get("iscrowd", 0) for obj in annotations], dtype=np.int64) + + boxes = [obj["bbox"] for obj in annotations] + # guard against no boxes via resizing + boxes = np.asarray(boxes, dtype=np.float32).reshape(-1, 4) + boxes[:, 2:] += boxes[:, :2] + boxes[:, 0::2] = boxes[:, 0::2].clip(min=0, max=image_width) + boxes[:, 1::2] = boxes[:, 1::2].clip(min=0, max=image_height) + + keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0]) + + new_target = {} + new_target["image_id"] = image_id + new_target["class_labels"] = classes[keep] + new_target["boxes"] = boxes[keep] + new_target["area"] = area[keep] + new_target["iscrowd"] = iscrowd[keep] + new_target["orig_size"] = np.asarray([int(image_height), int(image_width)], dtype=np.int64) + + if annotations and "keypoints" in annotations[0]: + keypoints = [obj["keypoints"] for obj in annotations] + # Converting the filtered keypoints list to a numpy array + keypoints = np.asarray(keypoints, dtype=np.float32) + # Apply the keep mask here to filter the relevant annotations + keypoints = keypoints[keep] + num_keypoints = keypoints.shape[0] + keypoints = keypoints.reshape((-1, 3)) if num_keypoints else keypoints + new_target["keypoints"] = keypoints + + if return_segmentation_masks: + segmentation_masks = [obj["segmentation"] for obj in annotations] + masks = convert_coco_poly_to_mask(segmentation_masks, image_height, image_width) + new_target["masks"] = masks[keep] + + return new_target + + +def masks_to_boxes(masks: np.ndarray) -> np.ndarray: + """ + Compute the bounding boxes around the provided panoptic segmentation masks. + + Args: + masks: masks in format `[number_masks, height, width]` where N is the number of masks + + Returns: + boxes: bounding boxes in format `[number_masks, 4]` in xyxy format + """ + if masks.size == 0: + return np.zeros((0, 4)) + + h, w = masks.shape[-2:] + y = np.arange(0, h, dtype=np.float32) + x = np.arange(0, w, dtype=np.float32) + # see https://github.com/pytorch/pytorch/issues/50276 + y, x = np.meshgrid(y, x, indexing="ij") + + x_mask = masks * np.expand_dims(x, axis=0) + x_max = x_mask.reshape(x_mask.shape[0], -1).max(-1) + x = np.ma.array(x_mask, mask=~(np.array(masks, dtype=bool))) + x_min = x.filled(fill_value=1e8) + x_min = x_min.reshape(x_min.shape[0], -1).min(-1) + + y_mask = masks * np.expand_dims(y, axis=0) + y_max = y_mask.reshape(x_mask.shape[0], -1).max(-1) + y = np.ma.array(y_mask, mask=~(np.array(masks, dtype=bool))) + y_min = y.filled(fill_value=1e8) + y_min = y_min.reshape(y_min.shape[0], -1).min(-1) + + return np.stack([x_min, y_min, x_max, y_max], 1) + + +# 2 functions below adapted from https://github.com/cocodataset/panopticapi/blob/master/panopticapi/utils.py +# Copyright (c) 2018, Alexander Kirillov +# All rights reserved. +def rgb_to_id(color): + """ + Converts RGB color to unique ID. + """ + if isinstance(color, np.ndarray) and len(color.shape) == 3: + if color.dtype == np.uint8: + color = color.astype(np.int32) + return color[:, :, 0] + 256 * color[:, :, 1] + 256 * 256 * color[:, :, 2] + return int(color[0] + 256 * color[1] + 256 * 256 * color[2]) + + +def prepare_coco_panoptic_annotation( + image: np.ndarray, + target: dict, + masks_path: str | pathlib.Path, + return_masks: bool = True, + input_data_format: ChannelDimension | str = None, +) -> dict: + """ + Prepare a coco panoptic annotation for CONDITIONAL_DETR. + """ + image_height, image_width = get_image_size(image, channel_dim=input_data_format) + annotation_path = pathlib.Path(masks_path) / target["file_name"] + + new_target = {} + new_target["image_id"] = np.asarray([target["image_id"] if "image_id" in target else target["id"]], dtype=np.int64) + new_target["size"] = np.asarray([image_height, image_width], dtype=np.int64) + new_target["orig_size"] = np.asarray([image_height, image_width], dtype=np.int64) + + if "segments_info" in target: + masks = np.asarray(PIL.Image.open(annotation_path), dtype=np.uint32) + masks = rgb_to_id(masks) + + ids = np.array([segment_info["id"] for segment_info in target["segments_info"]]) + masks = masks == ids[:, None, None] + masks = masks.astype(np.uint8) + if return_masks: + new_target["masks"] = masks + new_target["boxes"] = masks_to_boxes(masks) + new_target["class_labels"] = np.array( + [segment_info["category_id"] for segment_info in target["segments_info"]], dtype=np.int64 + ) + new_target["iscrowd"] = np.asarray( + [segment_info["iscrowd"] for segment_info in target["segments_info"]], dtype=np.int64 + ) + new_target["area"] = np.asarray( + [segment_info["area"] for segment_info in target["segments_info"]], dtype=np.float32 + ) + + return new_target + + +# Adapted from transformers.models.conditional_detr.image_processing_conditional_detr.binary_mask_to_rle +def binary_mask_to_rle(mask): + """ + Converts given binary mask of shape `(height, width)` to the run-length encoding (RLE) format. + + Args: + mask (`torch.Tensor` or `numpy.array`): + A binary mask tensor of shape `(height, width)` where 0 denotes background and 1 denotes the target + segment_id or class_id. + Returns: + `List`: Run-length encoded list of the binary mask. Refer to COCO API for more information about the RLE + format. + """ + from ...utils import is_torch_tensor + + if is_torch_tensor(mask): + mask = mask.numpy() + + pixels = mask.flatten() + pixels = np.concatenate([[0], pixels, [0]]) + runs = np.where(pixels[1:] != pixels[:-1])[0] + 1 + runs[1::2] -= runs[::2] + return list(runs) + + +# Adapted from transformers.models.conditional_detr.image_processing_conditional_detr.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_area = (mask_probs[k] >= mask_threshold).sum() + mask_exists = mask_k_area > 0 and original_area > 0 + + # Eliminate disconnected tiny segments + 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, mask_k + + +# Adapted from transformers.models.conditional_detr.image_processing_conditional_detr.compute_segments +def compute_segments( + mask_probs, + pred_scores, + pred_labels, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + label_ids_to_fuse: set[int] | None = None, + target_size: tuple[int, int] | None = None, +): + import torch + from torch import nn + + 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.int32, device=mask_probs.device) + segments: list[dict] = [] + + if target_size is not None: + mask_probs = nn.functional.interpolate( + mask_probs.unsqueeze(0), size=target_size, mode="bilinear", align_corners=False + )[0] + + current_segment_id = 0 + + # Weigh each mask by its prediction score + mask_probs *= pred_scores.view(-1, 1, 1) + mask_labels = mask_probs.argmax(0) # [height, width] + + # Keep track of instances of each class + stuff_memory_list: dict[str, int] = {} + for k in range(pred_labels.shape[0]): + pred_class = pred_labels[k].item() + should_fuse = pred_class in label_ids_to_fuse + + # Check if mask exists and large enough to be a segment + mask_exists, mask_k = check_segment_validity( + mask_labels, mask_probs, k, mask_threshold, overlap_mask_area_threshold + ) + + if mask_exists: + if pred_class in stuff_memory_list: + current_segment_id = stuff_memory_list[pred_class] + else: + current_segment_id += 1 + + # Add current object segment to final segmentation map + segmentation[mask_k] = current_segment_id + segment_score = round(pred_scores[k].item(), 6) + segments.append( + { + "id": current_segment_id, + "label_id": pred_class, + "was_fused": should_fuse, + "score": segment_score, + } + ) + if should_fuse: + stuff_memory_list[pred_class] = current_segment_id + + return segmentation, segments + + +# Adapted from transformers.models.conditional_detr.image_processing_conditional_detr.convert_segmentation_to_rle +def convert_segmentation_to_rle(segmentation): + """ + Converts given segmentation map of shape `(height, width)` to the run-length encoding (RLE) format. + + Args: + segmentation (`torch.Tensor` or `numpy.array`): + A segmentation map of shape `(height, width)` where each value denotes a segment or class id. + Returns: + `list[List]`: A list of lists, where each list is the run-length encoding of a segment / class id. + """ + import torch + + segment_ids = torch.unique(segmentation) + + run_length_encodings = [] + for idx in segment_ids: + mask = torch.where(segmentation == idx, 1, 0) + rle = binary_mask_to_rle(mask) + run_length_encodings.append(rle) + + return run_length_encodings + + +# Adapted from transformers.models.conditional_detr.image_processing_conditional_detr.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 +class ConditionalDetrImageProcessorPil(PilBackend): + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_DEFAULT_MEAN + image_std = IMAGENET_DEFAULT_STD + format = AnnotationFormat.COCO_DETECTION + do_resize = True + do_rescale = True + do_normalize = True + do_pad = True + size = {"shortest_edge": 800, "longest_edge": 1333} + default_to_square = False + model_input_names = ["pixel_values", "pixel_mask"] + valid_kwargs = ConditionalDetrImageProcessorKwargs + + def __init__(self, **kwargs: Unpack[ConditionalDetrImageProcessorKwargs]) -> None: + kwargs.setdefault("do_pad", kwargs.pop("pad_and_return_pixel_mask", self.do_pad)) + + size = kwargs.pop("size", None) + max_size = None if size is None else kwargs.pop("max_size", 1333) + size = size if size is not None else {"shortest_edge": 800, "longest_edge": 1333} + # Convert size dict for backwards compat with max_size parameter + if size is not None: + from ...image_processing_utils import get_size_dict + + kwargs["size"] = get_size_dict(size, max_size=max_size, default_to_square=False) + + # Backwards compatibility + do_convert_annotations = kwargs.get("do_convert_annotations") + do_normalize = kwargs.get("do_normalize") + if do_convert_annotations is None and getattr(self, "do_convert_annotations", None) is None: + self.do_convert_annotations = do_normalize if do_normalize is not None else self.do_normalize + + super().__init__(**kwargs) + + def prepare_annotation( + self, + image: np.ndarray, + target: dict, + format: AnnotationFormat | None = None, + return_segmentation_masks: bool | None = None, + masks_path: str | pathlib.Path | None = None, + input_data_format: str | ChannelDimension | None = None, + ) -> dict: + """ + Prepare an annotation for feeding into CONDITIONAL_DETR model. + """ + format = format if format is not None else self.format + + if format == AnnotationFormat.COCO_DETECTION: + return_segmentation_masks = False if return_segmentation_masks is None else return_segmentation_masks + target = prepare_coco_detection_annotation( + image, target, return_segmentation_masks, input_data_format=input_data_format + ) + elif format == AnnotationFormat.COCO_PANOPTIC: + return_segmentation_masks = True if return_segmentation_masks is None else return_segmentation_masks + target = prepare_coco_panoptic_annotation( + image, + target, + masks_path=masks_path, + return_masks=return_segmentation_masks, + input_data_format=input_data_format, + ) + else: + raise ValueError(f"Format {format} is not supported.") + return target + + def resize( + self, + image: np.ndarray, + size: SizeDict, + resample: Optional["PILImageResampling"] = None, + **kwargs, + ) -> np.ndarray: + """ + Resize the image to the given size. Size can be `min_size` (scalar) or `(height, width)` tuple. If size is an + int, smaller edge of the image will be matched to this number. + + Args: + image (`np.ndarray`): + Image to resize. + size (`SizeDict`): + Size of the image's `(height, width)` dimensions after resizing. Available options are: + - `{"height": int, "width": int}`: The image will be resized to the exact size `(height, width)`. + Do NOT keep the aspect ratio. + - `{"shortest_edge": int, "longest_edge": int}`: The image will be resized to a maximum size respecting + the aspect ratio and keeping the shortest edge less or equal to `shortest_edge` and the longest edge + less or equal to `longest_edge`. + - `{"max_height": int, "max_width": int}`: The image will be resized to the maximum size respecting the + aspect ratio and keeping the height less or equal to `max_height` and the width less or equal to + `max_width`. + resample (`PILImageResampling`, *optional*, defaults to `PILImageResampling.BILINEAR`): + Resampling filter to use if resizing the image. + """ + resample = resample if resample is not None else self.resample + + if size.shortest_edge and size.longest_edge: + # Resize the image so that the shortest edge or the longest edge is of the given size + # while maintaining the aspect ratio of the original image. + new_size = get_size_with_aspect_ratio( + image.shape[-2:], + size.shortest_edge, + size.longest_edge or size.shortest_edge, + ) + elif size.max_height and size.max_width: + new_size = get_image_size_for_max_height_width(image.shape[-2:], size.max_height, size.max_width) + elif size.height and size.width: + new_size = (size.height, size.width) + else: + raise ValueError( + f"Size must contain 'height' and 'width' keys or 'shortest_edge' and 'longest_edge' keys. Got {size}." + ) + + image = super().resize( + image, + size=SizeDict(height=new_size[0], width=new_size[1]), + resample=resample, + **kwargs, + ) + return image + + def resize_annotation( + self, + annotation: dict[str, Any], + orig_size: tuple[int, int], + target_size: tuple[int, int], + threshold: float = 0.5, + resample: Optional["PILImageResampling"] = PILImageResampling.NEAREST, + ): + """ + Resizes an annotation to a target size. + + Args: + annotation (`dict[str, Any]`): + The annotation dictionary. + orig_size (`tuple[int, int]`): + The original size of the input image. + target_size (`tuple[int, int]`): + The target size of the image, as returned by the preprocessing `resize` step. + threshold (`float`, *optional*, defaults to 0.5): + The threshold used to binarize the segmentation masks. + resample (`PILImageResampling`, defaults to `PILImageResampling.NEAREST`): + The resampling filter to use when resizing the masks. + """ + ratios = tuple(float(s) / float(s_orig) for s, s_orig in zip(target_size, orig_size)) + ratio_height, ratio_width = ratios + + new_annotation = {} + new_annotation["size"] = target_size + + for key, value in annotation.items(): + if key == "boxes": + boxes = value + scaled_boxes = boxes * np.asarray( + [ratio_width, ratio_height, ratio_width, ratio_height], dtype=np.float32 + ) + new_annotation["boxes"] = scaled_boxes + elif key == "area": + area = value + scaled_area = area * (ratio_width * ratio_height) + new_annotation["area"] = scaled_area + elif key == "masks": + masks = value[:, None] + masks = np.array([resize(mask, target_size, resample=resample) for mask in masks]) + masks = masks.astype(np.float32) + masks = masks[:, 0] > threshold + new_annotation["masks"] = masks + elif key == "size": + new_annotation["size"] = target_size + else: + new_annotation[key] = value + + return new_annotation + + def normalize_annotation(self, annotation: dict, image_size: tuple[int, int]) -> dict: + image_height, image_width = image_size + norm_annotation = {} + for key, value in annotation.items(): + if key == "boxes": + boxes = value + boxes = corners_to_center_format(boxes) + boxes /= np.asarray([image_width, image_height, image_width, image_height], dtype=np.float32) + norm_annotation[key] = boxes + else: + norm_annotation[key] = value + return norm_annotation + + def _update_annotation_for_padded_image( + self, + annotation: dict, + input_image_size: tuple[int, int], + output_image_size: tuple[int, int], + padding, + update_bboxes, + ) -> dict: + """ + Update the annotation for a padded image. + """ + new_annotation = {} + new_annotation["size"] = output_image_size + ratio_height, ratio_width = (input / output for output, input in zip(output_image_size, input_image_size)) + + for key, value in annotation.items(): + if key == "masks": + masks = value + masks = pad( + masks, + padding, + mode=PaddingMode.CONSTANT, + constant_values=0, + input_data_format=ChannelDimension.FIRST, + ) + masks = safe_squeeze(masks, 1) + new_annotation["masks"] = masks + elif key == "boxes" and update_bboxes: + boxes = value + boxes *= np.asarray( + [ + input_image_size[1] / output_image_size[1], + input_image_size[0] / output_image_size[0], + input_image_size[1] / output_image_size[1], + input_image_size[0] / output_image_size[0], + ] + ) + new_annotation["boxes"] = boxes + elif key == "size": + new_annotation["size"] = output_image_size + else: + new_annotation[key] = value + return new_annotation + + def pad( + self, + image: np.ndarray, + padded_size: tuple[int, int], + annotation: dict[str, Any] | None = None, + update_bboxes: bool = True, + fill: int = 0, + ): + input_height, input_width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + output_height, output_width = padded_size + padding_bottom = output_height - input_height + padding_right = output_width - input_width + if padding_bottom < 0 or padding_right < 0: + raise ValueError( + f"Padding dimensions are negative. Please make sure that the padded size is larger than the " + f"original size. Got padded size: {padded_size}, original size: {(input_height, input_width)}." + ) + if (input_height, input_width) != padded_size: + padding = ((0, padding_bottom), (0, padding_right)) + image = pad( + image, + padding, + mode=PaddingMode.CONSTANT, + constant_values=fill, + data_format=ChannelDimension.FIRST, + input_data_format=ChannelDimension.FIRST, + ) + if annotation is not None: + annotation = self._update_annotation_for_padded_image( + annotation, (input_height, input_width), (output_height, output_width), padding, update_bboxes + ) + + # Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding. + pixel_mask = np.zeros(padded_size, dtype=np.int64) + pixel_mask[:input_height, :input_width] = 1 + + return image, pixel_mask, annotation + + @auto_docstring + def preprocess( + self, + images: ImageInput, + annotations: AnnotationType | list[AnnotationType] | None = None, + return_segmentation_masks: bool | None = None, + masks_path: str | pathlib.Path | None = None, + **kwargs: Unpack[ConditionalDetrImageProcessorKwargs], + ) -> BatchFeature: + r""" + annotations (`AnnotationType` or `list[AnnotationType]`, *optional*): + Annotations to transform according to the padding that is applied to the images. + return_segmentation_masks (`bool`, *optional*, defaults to `self.return_segmentation_masks`): + Whether to return segmentation masks. + masks_path (`str` or `pathlib.Path`, *optional*): + Path to the directory containing the segmentation masks. + """ + return super().preprocess(images, annotations, return_segmentation_masks, masks_path, **kwargs) + + def _preprocess( + self, + images: list[np.ndarray], + annotations: AnnotationType | list[AnnotationType] | None, + return_segmentation_masks: bool, + masks_path: str | pathlib.Path | None, + do_resize: bool, + size: SizeDict, + resample: "PILImageResampling | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + do_convert_annotations: bool, + image_mean: float | list[float] | None, + image_std: float | list[float] | None, + do_pad: bool, + pad_size: SizeDict | None, + format: str | AnnotationFormat | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + """ + Preprocess an image or a batch of images so that it can be used by the model. + """ + if annotations is not None and isinstance(annotations, dict): + annotations = [annotations] + + if annotations is not None and len(images) != len(annotations): + raise ValueError( + f"The number of images ({len(images)}) and annotations ({len(annotations)}) do not match." + ) + + format = AnnotationFormat(format) + if annotations is not None: + validate_annotations(format, SUPPORTED_ANNOTATION_FORMATS, annotations) + + if ( + masks_path is not None + and format == AnnotationFormat.COCO_PANOPTIC + and not isinstance(masks_path, (pathlib.Path, str)) + ): + raise ValueError( + "The path to the directory containing the mask PNG files should be provided as a" + f" `pathlib.Path` or string object, but is {type(masks_path)} instead." + ) + + data = {} + + # Import torch if needed for tensor conversion + if return_tensors == "pt": + if not is_torch_available(): + raise ImportError("PyTorch is required for tensor conversion.") + + processed_images = [] + processed_annotations = [] + pixel_masks = [] # Initialize pixel_masks here + for image, annotation in zip(images, annotations if annotations is not None else [None] * len(images)): + # prepare (COCO annotations as a list of Dict -> CONDITIONAL_DETR target as a single Dict per image) + if annotations is not None: + annotation = self.prepare_annotation( + image, + annotation, + format, + return_segmentation_masks=return_segmentation_masks, + masks_path=masks_path, + input_data_format=ChannelDimension.FIRST, + ) + + if do_resize: + resized_image = self.resize(image, size=size, resample=resample) + if annotations is not None: + annotation = self.resize_annotation( + annotation, + orig_size=get_image_size(image, channel_dim=ChannelDimension.FIRST), + target_size=get_image_size(resized_image, channel_dim=ChannelDimension.FIRST), + ) + image = resized_image + + if do_rescale: + image = self.rescale(image, rescale_factor) + if do_normalize: + image = self.normalize(image, image_mean, image_std) + + if do_convert_annotations and annotations is not None: + annotation = self.normalize_annotation(annotation, get_image_size(image, ChannelDimension.FIRST)) + + processed_images.append(image) + processed_annotations.append(annotation) + images = processed_images + annotations = processed_annotations if annotations is not None else None + + if do_pad: + # depends on all resized image shapes so we need another loop + if pad_size is not None: + padded_size = (pad_size.height, pad_size.width) + else: + padded_size = get_max_height_width(images, input_data_format=ChannelDimension.FIRST) + + padded_images = [] + padded_annotations = [] + for image, annotation in zip(images, annotations if annotations is not None else [None] * len(images)): + # Pads images and returns their mask: {'pixel_values': ..., 'pixel_mask': ...} + image_height, image_width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + if padded_size == (image_height, image_width): + padded_images.append(image) + pixel_masks.append(np.ones(padded_size, dtype=np.int64)) + padded_annotations.append(annotation) + continue + image, pixel_mask, annotation = self.pad( + image, padded_size, annotation=annotation, update_bboxes=do_convert_annotations + ) + padded_images.append(image) + padded_annotations.append(annotation) + pixel_masks.append(pixel_mask) + images = padded_images + annotations = padded_annotations if annotations is not None else None + data.update({"pixel_mask": pixel_masks}) + + data.update({"pixel_values": images}) + encoded_inputs = BatchFeature(data, tensor_type=return_tensors) + if annotations is not None: + encoded_inputs["labels"] = [ + BatchFeature(annotation, tensor_type=return_tensors) for annotation in annotations + ] + return encoded_inputs + + @requires(backends=("torch",)) + def post_process_object_detection( + self, outputs, threshold: float = 0.5, target_sizes: TensorType | list[tuple] = None, top_k: int = 100 + ): + """ + Converts the raw output of [`ConditionalDetrForObjectDetection`] into final bounding boxes in (top_left_x, + top_left_y, bottom_right_x, bottom_right_y) format. Only supports PyTorch. + + Args: + outputs ([`ConditionalDetrObjectDetectionOutput`]): + Raw outputs of the model. + threshold (`float`, *optional*): + Score threshold to keep object detection predictions. + target_sizes (`torch.Tensor` or `list[tuple[int, int]]`, *optional*): + Tensor of shape `(batch_size, 2)` or list of tuples (`tuple[int, int]`) containing the target size + (height, width) of each image in the batch. If left to None, predictions will not be resized. + top_k (`int`, *optional*, defaults to 100): + Keep only top k bounding boxes before filtering by thresholding. + + Returns: + `list[Dict]`: A list of dictionaries, each dictionary containing the scores, labels and boxes for an image + in the batch as predicted by the model. + """ + requires_backends(self, ["torch"]) + out_logits, out_bbox = outputs.logits, outputs.pred_boxes + + if target_sizes is not None: + if len(out_logits) != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + prob = out_logits.sigmoid() + prob = prob.view(out_logits.shape[0], -1) + k_value = min(top_k, prob.size(1)) + topk_values, topk_indexes = torch.topk(prob, k_value, dim=1) + scores = topk_values + topk_boxes = torch.div(topk_indexes, out_logits.shape[2], rounding_mode="floor") + labels = topk_indexes % out_logits.shape[2] + boxes = center_to_corners_format(out_bbox) + boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4)) + + # and from relative [0, 1] to absolute [0, height] coordinates + if target_sizes is not None: + if isinstance(target_sizes, list): + img_h = torch.Tensor([i[0] for i in target_sizes]) + img_w = torch.Tensor([i[1] for i in target_sizes]) + else: + img_h, img_w = target_sizes.unbind(1) + scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(boxes.device) + boxes = boxes * scale_fct[:, None, :] + + results = [] + for s, l, b in zip(scores, labels, boxes): + score = s[s > threshold] + label = l[s > threshold] + box = b[s > threshold] + results.append({"scores": score, "labels": label, "boxes": box}) + + return results + + @requires(backends=("torch",)) + def post_process_semantic_segmentation(self, outputs, target_sizes: list[tuple[int, int]] | None = None): + """ + Converts the output of [`ConditionalDetrForSegmentation`] into semantic segmentation maps. Only supports PyTorch. + + Args: + outputs ([`ConditionalDetrForSegmentation`]): + Raw outputs of the model. + target_sizes (`list[tuple[int, int]]`, *optional*): + A list of tuples (`tuple[int, int]`) containing the target size (height, width) of each image in the + batch. If unset, predictions will not be resized. + Returns: + `list[torch.Tensor]`: + A list of length `batch_size`, where each item is a semantic segmentation map of shape (height, width) + corresponding to the target_sizes entry (if `target_sizes` is specified). Each entry of each + `torch.Tensor` correspond to a semantic class id. + """ + requires_backends(self, ["torch"]) + class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes] + masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width] + + # Conditional DETR does not have a null class, so we use all classes + masks_classes = class_queries_logits.softmax(dim=-1) + masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # Semantic segmentation logits of shape (batch_size, num_classes, height, width) + segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs) + batch_size = class_queries_logits.shape[0] + + # Resize logits and compute semantic segmentation maps + if target_sizes is not None: + if batch_size != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + semantic_segmentation = [] + for idx in range(batch_size): + resized_logits = nn.functional.interpolate( + segmentation[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False + ) + semantic_map = resized_logits[0].argmax(dim=0) + semantic_segmentation.append(semantic_map) + else: + semantic_segmentation = segmentation.argmax(dim=1) + semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])] + + return semantic_segmentation + + @requires(backends=("torch",)) + def post_process_instance_segmentation( + self, + outputs, + threshold: float = 0.5, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + target_sizes: list[tuple[int, int]] | None = None, + return_coco_annotation: bool | None = False, + ) -> list[dict]: + """ + Converts the output of [`ConditionalDetrForSegmentation`] into instance segmentation predictions. Only supports PyTorch. + + Args: + outputs ([`ConditionalDetrForSegmentation`]): + Raw outputs of the model. + threshold (`float`, *optional*, defaults to 0.5): + The probability score threshold to keep predicted instance masks. + mask_threshold (`float`, *optional*, defaults to 0.5): + Threshold to use when turning the predicted masks into binary values. + overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8): + The overlap mask area threshold to merge or discard small disconnected parts within each binary + instance mask. + target_sizes (`list[Tuple]`, *optional*): + List of length (batch_size), where each list item (`tuple[int, int]]`) corresponds to the requested + final size (height, width) of each prediction. If unset, predictions will not be resized. + return_coco_annotation (`bool`, *optional*): + Defaults to `False`. If set to `True`, segmentation maps are returned in COCO run-length encoding (RLE) + format. + Returns: + `list[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys: + - **segmentation** -- A tensor of shape `(height, width)` where each pixel represents a `segment_id` or + `list[List]` run-length encoding (RLE) of the segmentation map if return_coco_annotation is set to + `True`. Set to `None` if no mask if found above `threshold`. + - **segments_info** -- A dictionary that contains additional information on each segment. + - **id** -- An integer representing the `segment_id`. + - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`. + - **score** -- Prediction score of segment with `segment_id`. + """ + if not is_torch_available(): + raise ImportError("PyTorch is required for post-processing.") + import torch + from torch import nn + + class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width] + + batch_size = class_queries_logits.shape[0] + num_labels = class_queries_logits.shape[-1] - 1 + + mask_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # Predicted label and score of each query (batch_size, num_queries) + pred_scores, pred_labels = nn.functional.softmax(class_queries_logits, dim=-1).max(-1) + + # Loop over items in batch size + results: list[dict[str, TensorType]] = [] + + for i in range(batch_size): + mask_probs_item, pred_scores_item, pred_labels_item = remove_low_and_no_objects( + mask_probs[i], pred_scores[i], pred_labels[i], threshold, num_labels + ) + + # No mask found + if mask_probs_item.shape[0] <= 0: + height, width = target_sizes[i] if target_sizes is not None else mask_probs_item.shape[1:] + segmentation = torch.zeros((height, width)) - 1 + results.append({"segmentation": segmentation, "segments_info": []}) + continue + + # Get segmentation map and segment information of batch item + target_size = target_sizes[i] if target_sizes is not None else None + segmentation, segments = compute_segments( + mask_probs=mask_probs_item, + pred_scores=pred_scores_item, + pred_labels=pred_labels_item, + mask_threshold=mask_threshold, + overlap_mask_area_threshold=overlap_mask_area_threshold, + label_ids_to_fuse=[], + target_size=target_size, + ) + + # Return segmentation map in run-length encoding (RLE) format + if return_coco_annotation: + segmentation = convert_segmentation_to_rle(segmentation) + + results.append({"segmentation": segmentation, "segments_info": segments}) + return results + + @requires(backends=("torch",)) + def post_process_panoptic_segmentation( + self, + outputs, + threshold: float = 0.5, + mask_threshold: float = 0.5, + overlap_mask_area_threshold: float = 0.8, + label_ids_to_fuse: set[int] | None = None, + target_sizes: list[tuple[int, int]] | None = None, + ) -> list[dict]: + """ + Converts the output of [`ConditionalDetrForSegmentation`] into image panoptic segmentation predictions. Only supports + PyTorch. + + Args: + outputs ([`ConditionalDetrForSegmentation`]): + The outputs from [`ConditionalDetrForSegmentation`]. + threshold (`float`, *optional*, defaults to 0.5): + The probability score threshold to keep predicted instance masks. + mask_threshold (`float`, *optional*, defaults to 0.5): + Threshold to use when turning the predicted masks into binary values. + overlap_mask_area_threshold (`float`, *optional*, defaults to 0.8): + The overlap mask area threshold to merge or discard small disconnected parts within each binary + instance mask. + label_ids_to_fuse (`Set[int]`, *optional*): + The labels in this state will have all their instances be fused together. For instance we could say + there can only be one sky in an image, but several persons, so the label ID for sky would be in that + set, but not the one for person. + target_sizes (`list[Tuple]`, *optional*): + List of length (batch_size), where each list item (`tuple[int, int]]`) corresponds to the requested + final size (height, width) of each prediction in batch. If unset, predictions will not be resized. + Returns: + `list[Dict]`: A list of dictionaries, one per image, each dictionary containing two keys: + - **segmentation** -- a tensor of shape `(height, width)` where each pixel represents a `segment_id` or + `None` if no mask if found above `threshold`. If `target_sizes` is specified, segmentation is resized to + the corresponding `target_sizes` entry. + - **segments_info** -- A dictionary that contains additional information on each segment. + - **id** -- an integer representing the `segment_id`. + - **label_id** -- An integer representing the label / semantic class id corresponding to `segment_id`. + - **was_fused** -- a boolean, `True` if `label_id` was in `label_ids_to_fuse`, `False` otherwise. + Multiple instances of the same class / label were fused and assigned a single `segment_id`. + - **score** -- Prediction score of segment with `segment_id`. + """ + + if label_ids_to_fuse is None: + logger.warning_once("`label_ids_to_fuse` unset. No instance will be fused.") + label_ids_to_fuse = set() + + if not is_torch_available(): + raise ImportError("PyTorch is required for post-processing.") + import torch + from torch import nn + + class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes+1] + masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width] + + batch_size = class_queries_logits.shape[0] + num_labels = class_queries_logits.shape[-1] - 1 + + mask_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # Predicted label and score of each query (batch_size, num_queries) + pred_scores, pred_labels = nn.functional.softmax(class_queries_logits, dim=-1).max(-1) + + # Loop over items in batch size + results: list[dict[str, TensorType]] = [] + + for i in range(batch_size): + mask_probs_item, pred_scores_item, pred_labels_item = remove_low_and_no_objects( + mask_probs[i], pred_scores[i], pred_labels[i], threshold, num_labels + ) + + # No mask found + if mask_probs_item.shape[0] <= 0: + height, width = target_sizes[i] if target_sizes is not None else mask_probs_item.shape[1:] + segmentation = torch.zeros((height, width)) - 1 + results.append({"segmentation": segmentation, "segments_info": []}) + continue + + # Get segmentation map and segment information of batch item + target_size = target_sizes[i] if target_sizes is not None else None + segmentation, segments = compute_segments( + mask_probs=mask_probs_item, + pred_scores=pred_scores_item, + pred_labels=pred_labels_item, + mask_threshold=mask_threshold, + overlap_mask_area_threshold=overlap_mask_area_threshold, + label_ids_to_fuse=label_ids_to_fuse, + target_size=target_size, + ) + + results.append({"segmentation": segmentation, "segments_info": segments}) + return results + + +__all__ = ["ConditionalDetrImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/conditional_detr/modeling_conditional_detr.py b/third_party/transformers/src/transformers/models/conditional_detr/modeling_conditional_detr.py new file mode 100644 index 0000000000000000000000000000000000000000..2f10c81b38e1040ab88e5c1dfd756bbfe0004dfa --- /dev/null +++ b/third_party/transformers/src/transformers/models/conditional_detr/modeling_conditional_detr.py @@ -0,0 +1,1827 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/conditional_detr/modular_conditional_detr.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_conditional_detr.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2022 Microsoft Research Asia 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 math +from collections.abc import Callable +from dataclasses import dataclass + +import torch +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...backbone_utils import load_backbone +from ...masking_utils import create_bidirectional_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithCrossAttentions, Seq2SeqModelOutput +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...pytorch_utils import compile_compatible_method_lru_cache +from ...utils import ModelOutput, TransformersKwargs, auto_docstring +from ...utils.generic import can_return_tuple, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from .configuration_conditional_detr import ConditionalDetrConfig + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of the CONDITIONAL_DETR decoder. This class adds one attribute to BaseModelOutputWithCrossAttentions, + namely an optional stack of intermediate decoder activations, i.e. the output of each decoder layer, each of them + gone through a layernorm. This is useful when training the model with auxiliary decoding losses. + """ +) +class ConditionalDetrDecoderOutput(BaseModelOutputWithCrossAttentions): + r""" + 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. + intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`): + Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a + layernorm. + reference_points (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, 2 (anchor points))`): + Reference points (reference points of each layer of the decoder). + """ + + intermediate_hidden_states: torch.FloatTensor | None = None + + reference_points: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of the CONDITIONAL_DETR encoder-decoder model. This class adds one attribute to Seq2SeqModelOutput, + namely an optional stack of intermediate decoder activations, i.e. the output of each decoder layer, each of them + gone through a layernorm. This is useful when training the model with auxiliary decoding losses. + """ +) +class ConditionalDetrModelOutput(Seq2SeqModelOutput): + 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. + intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, sequence_length, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`): + Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a + layernorm. + reference_points (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, 2 (anchor points))`): + Reference points (reference points of each layer of the decoder). + """ + + intermediate_hidden_states: torch.FloatTensor | None = None + + reference_points: tuple[torch.FloatTensor] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Output type of [`ConditionalDetrForObjectDetection`]. + """ +) +class ConditionalDetrObjectDetectionOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)): + Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a + bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized + scale-invariant IoU loss. + loss_dict (`Dict`, *optional*): + A dictionary containing the individual losses. Useful for logging. + logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`): + Classification logits (including no-object) for all queries. + pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): + Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These + values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding + possible padding). You can use [`~ConditionalDetrImageProcessor.post_process_object_detection`] to retrieve the + unnormalized bounding boxes. + auxiliary_outputs (`list[Dict]`, *optional*): + Optional, only returned when auxiliary losses are activated (i.e. `config.auxiliary_loss` is set to `True`) + and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and + `pred_boxes`) for each decoder layer. + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the decoder of the model. + """ + + loss: torch.FloatTensor | None = None + loss_dict: dict | None = None + logits: torch.FloatTensor | None = None + pred_boxes: torch.FloatTensor | None = None + auxiliary_outputs: list[dict] | None = None + last_hidden_state: torch.FloatTensor | 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 + + +@dataclass +@auto_docstring( + custom_intro=""" + Output type of [`ConditionalDetrForSegmentation`]. + """ +) +class ConditionalDetrSegmentationOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)): + Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a + bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized + scale-invariant IoU loss. + loss_dict (`Dict`, *optional*): + A dictionary containing the individual losses. Useful for logging. + logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`): + Classification logits (including no-object) for all queries. + pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`): + Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These + values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding + possible padding). You can use [`~ConditionalDetrImageProcessor.post_process_object_detection`] to retrieve the + unnormalized bounding boxes. + pred_masks (`torch.FloatTensor` of shape `(batch_size, num_queries, height/4, width/4)`): + Segmentation masks logits for all queries. See also + [`~ConditionalDetrImageProcessor.post_process_semantic_segmentation`] or + [`~ConditionalDetrImageProcessor.post_process_instance_segmentation`] + [`~ConditionalDetrImageProcessor.post_process_panoptic_segmentation`] to evaluate semantic, instance and panoptic + segmentation masks respectively. + auxiliary_outputs (`list[Dict]`, *optional*): + Optional, only returned when auxiliary losses are activated (i.e. `config.auxiliary_loss` is set to `True`) + and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and + `pred_boxes`) for each decoder layer. + last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Sequence of hidden-states at the output of the last layer of the decoder of the model. + """ + + loss: torch.FloatTensor | None = None + loss_dict: dict | None = None + logits: torch.FloatTensor | None = None + pred_boxes: torch.FloatTensor | None = None + pred_masks: torch.FloatTensor | None = None + auxiliary_outputs: list[dict] | None = None + last_hidden_state: torch.FloatTensor | 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 + + +class ConditionalDetrFrozenBatchNorm2d(nn.Module): + """ + BatchNorm2d where the batch statistics and the affine parameters are fixed. + + Copy-paste from torchvision.misc.ops with added eps before rqsrt, without which any other models than + torchvision.models.resnet[18,34,50,101] produce nans. + """ + + def __init__(self, n): + super().__init__() + self.register_buffer("weight", torch.ones(n)) + self.register_buffer("bias", torch.zeros(n)) + self.register_buffer("running_mean", torch.zeros(n)) + self.register_buffer("running_var", torch.ones(n)) + + def _load_from_state_dict( + self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ): + num_batches_tracked_key = prefix + "num_batches_tracked" + if num_batches_tracked_key in state_dict: + del state_dict[num_batches_tracked_key] + + super()._load_from_state_dict( + state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs + ) + + def forward(self, x): + # move reshapes to the beginning + # to make it user-friendly + weight = self.weight.reshape(1, -1, 1, 1) + bias = self.bias.reshape(1, -1, 1, 1) + running_var = self.running_var.reshape(1, -1, 1, 1) + running_mean = self.running_mean.reshape(1, -1, 1, 1) + epsilon = 1e-5 + scale = weight * (running_var + epsilon).rsqrt() + bias = bias - running_mean * scale + return x * scale + bias + + +def replace_batch_norm(model): + r""" + Recursively replace all `torch.nn.BatchNorm2d` with `ConditionalDetrFrozenBatchNorm2d`. + + Args: + model (torch.nn.Module): + input model + """ + for name, module in model.named_children(): + if isinstance(module, nn.BatchNorm2d): + new_module = ConditionalDetrFrozenBatchNorm2d(module.num_features) + + if module.weight.device != torch.device("meta"): + new_module.weight.copy_(module.weight) + new_module.bias.copy_(module.bias) + new_module.running_mean.copy_(module.running_mean) + new_module.running_var.copy_(module.running_var) + + model._modules[name] = new_module + + if len(list(module.children())) > 0: + replace_batch_norm(module) + + +class ConditionalDetrConvEncoder(nn.Module): + """ + Convolutional backbone, using either the AutoBackbone API or one from the timm library. + + nn.BatchNorm2d layers are replaced by ConditionalDetrFrozenBatchNorm2d as defined above. + + """ + + def __init__(self, config): + super().__init__() + + self.config = config + + backbone = load_backbone(config) + self.intermediate_channel_sizes = backbone.channels + + # replace batch norm by frozen batch norm + with torch.no_grad(): + replace_batch_norm(backbone) + + # We used to load with timm library directly instead of the AutoBackbone API + # so we need to unwrap the `backbone._backbone` module to load weights without mismatch + is_timm_model = False + if hasattr(backbone, "_backbone"): + backbone = backbone._backbone + is_timm_model = True + self.model = backbone + + backbone_model_type = config.backbone_config.model_type + if "resnet" in backbone_model_type: + for name, parameter in self.model.named_parameters(): + if is_timm_model: + if "layer2" not in name and "layer3" not in name and "layer4" not in name: + parameter.requires_grad_(False) + else: + if "stage.1" not in name and "stage.2" not in name and "stage.3" not in name: + parameter.requires_grad_(False) + + def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor): + # send pixel_values through the model to get list of feature maps + features = self.model(pixel_values) + if isinstance(features, dict): + features = features.feature_maps + + out = [] + for feature_map in features: + # downsample pixel_mask to match shape of corresponding feature_map + mask = nn.functional.interpolate(pixel_mask[None].float(), size=feature_map.shape[-2:]).to(torch.bool)[0] + out.append((feature_map, mask)) + return out + + +class ConditionalDetrSinePositionEmbedding(nn.Module): + """ + This is a more standard version of the position embedding, very similar to the one used by the Attention is all you + need paper, generalized to work on images. + """ + + def __init__( + self, + num_position_features: int = 64, + temperature: int = 10000, + normalize: bool = False, + scale: float | None = None, + ): + super().__init__() + if scale is not None and normalize is False: + raise ValueError("normalize should be True if scale is passed") + self.num_position_features = num_position_features + self.temperature = temperature + self.normalize = normalize + self.scale = 2 * math.pi if scale is None else scale + + @compile_compatible_method_lru_cache(maxsize=1) + def forward( + self, + shape: torch.Size, + device: torch.device | str, + dtype: torch.dtype, + mask: torch.Tensor | None = None, + ) -> torch.Tensor: + if mask is None: + mask = torch.zeros((shape[0], shape[2], shape[3]), device=device, dtype=torch.bool) + y_embed = mask.cumsum(1, dtype=dtype) + x_embed = mask.cumsum(2, dtype=dtype) + if self.normalize: + eps = 1e-6 + y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale + x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale + + dim_t = torch.arange(self.num_position_features, dtype=torch.int64, device=device).to(dtype) + dim_t = self.temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / self.num_position_features) + + pos_x = x_embed[:, :, :, None] / dim_t + pos_y = y_embed[:, :, :, None] / dim_t + pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3) + pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3) + pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2) + # Flatten spatial dimensions and permute to (batch_size, sequence_length, hidden_size) format + # expected by the encoder + pos = pos.flatten(2).permute(0, 2, 1) + return pos + + +class ConditionalDetrLearnedPositionEmbedding(nn.Module): + """ + This module learns positional embeddings up to a fixed maximum size. + """ + + def __init__(self, embedding_dim=256): + super().__init__() + self.row_embeddings = nn.Embedding(50, embedding_dim) + self.column_embeddings = nn.Embedding(50, embedding_dim) + + @compile_compatible_method_lru_cache(maxsize=1) + def forward( + self, + shape: torch.Size, + device: torch.device | str, + dtype: torch.dtype, + mask: torch.Tensor | None = None, + ): + height, width = shape[-2:] + width_values = torch.arange(width, device=device) + height_values = torch.arange(height, device=device) + x_emb = self.column_embeddings(width_values) + y_emb = self.row_embeddings(height_values) + pos = torch.cat([x_emb.unsqueeze(0).repeat(height, 1, 1), y_emb.unsqueeze(1).repeat(1, width, 1)], dim=-1) + pos = pos.permute(2, 0, 1) + pos = pos.unsqueeze(0) + pos = pos.repeat(shape[0], 1, 1, 1) + # Flatten spatial dimensions and permute to (batch_size, sequence_length, hidden_size) format + # expected by the encoder + pos = pos.flatten(2).permute(0, 2, 1) + return pos + + +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 ConditionalDetrSelfAttention(nn.Module): + """ + Multi-headed self-attention from 'Attention Is All You Need' paper. + + In CONDITIONAL_DETR, position embeddings are added to both queries and keys (but not values) in self-attention. + """ + + def __init__( + self, + config: ConditionalDetrConfig, + hidden_size: int, + num_attention_heads: int, + dropout: float = 0.0, + bias: bool = True, + ): + super().__init__() + self.config = config + self.head_dim = hidden_size // num_attention_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = dropout + self.is_causal = False + + self.k_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.v_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.o_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + position_embeddings: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Position embeddings are added to both queries and keys (but not values). + """ + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_key_input = hidden_states + position_embeddings if position_embeddings is not None else hidden_states + + query_states = self.q_proj(query_key_input).view(hidden_shape).transpose(1, 2) + key_states = self.k_proj(query_key_input).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(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() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class ConditionalDetrDecoderSelfAttention(nn.Module): + """ + Multi-headed self-attention for Conditional DETR decoder layers. + + This attention module handles separate content and position projections, which are then combined + before applying standard self-attention. Position embeddings are added to both queries and keys. + """ + + def __init__( + self, + config: ConditionalDetrConfig, + hidden_size: int, + num_attention_heads: int, + dropout: float | int = 0.0, + ): + super().__init__() + self.config = config + self.hidden_size = hidden_size + self.head_dim = hidden_size // num_attention_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = dropout + self.is_causal = False + + # Content and position projections + self.q_content_proj = nn.Linear(hidden_size, hidden_size) + self.q_pos_proj = nn.Linear(hidden_size, hidden_size) + self.k_content_proj = nn.Linear(hidden_size, hidden_size) + self.k_pos_proj = nn.Linear(hidden_size, hidden_size) + self.v_proj = nn.Linear(hidden_size, hidden_size) + self.o_proj = nn.Linear(hidden_size, hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + query_position_embeddings: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`): + Input hidden states from the decoder layer. + query_position_embeddings (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`): + Position embeddings for queries and keys. Required (unlike standard attention). Processed through + separate position projections (`q_pos_proj`, `k_pos_proj`) and added to content projections. + attention_mask (`torch.Tensor` of shape `(batch_size, 1, num_queries, num_queries)`, *optional*): + Attention mask to avoid attending to padding tokens. + """ + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = ( + (self.q_content_proj(hidden_states) + self.q_pos_proj(query_position_embeddings)) + .view(hidden_shape) + .transpose(1, 2) + ) + key_states = ( + (self.k_content_proj(hidden_states) + self.k_pos_proj(query_position_embeddings)) + .view(hidden_shape) + .transpose(1, 2) + ) + value_states = self.v_proj(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() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class ConditionalDetrDecoderCrossAttention(nn.Module): + """ + Multi-headed cross-attention for Conditional DETR decoder layers. + + This attention module handles the special cross-attention logic in Conditional DETR: + - Separate content and position projections for queries and keys + - Concatenation of query sine embeddings with queries (doubling query dimension) + - Concatenation of key position embeddings with keys (doubling key dimension) + - Output dimension remains hidden_size despite doubled input dimensions + """ + + def __init__( + self, + config: ConditionalDetrConfig, + hidden_size: int, + num_attention_heads: int, + dropout: float | int = 0.0, + ): + super().__init__() + self.config = config + self.hidden_size = hidden_size + self.num_attention_heads = num_attention_heads + self.head_dim = hidden_size // num_attention_heads + self.attention_dropout = dropout + self.is_causal = False + + # Content and position projections + self.q_content_proj = nn.Linear(hidden_size, hidden_size) + self.q_pos_proj = nn.Linear(hidden_size, hidden_size) + self.k_content_proj = nn.Linear(hidden_size, hidden_size) + self.k_pos_proj = nn.Linear(hidden_size, hidden_size) + self.v_proj = nn.Linear(hidden_size, hidden_size) + self.q_pos_sine_proj = nn.Linear(hidden_size, hidden_size) + + # Output projection: input is hidden_size * 2 (from concatenated q/k), output is hidden_size + self.o_proj = nn.Linear(hidden_size, hidden_size) + + # Compute scaling for expanded head_dim (q and k have doubled dimensions after concatenation) + # This matches the original Conditional DETR implementation where embed_dim * 2 is used + expanded_head_dim = (hidden_size * 2) // num_attention_heads + self.scaling = expanded_head_dim**-0.5 + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + query_sine_embed: torch.Tensor, + encoder_position_embeddings: torch.Tensor, + query_position_embeddings: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`): + Decoder hidden states (queries). + encoder_hidden_states (`torch.Tensor` of shape `(batch_size, encoder_seq_len, hidden_size)`): + Encoder output hidden states (keys and values). + query_sine_embed (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`): + Sine position embeddings for queries. **Concatenated** (not added) with query content, + doubling the query dimension. + encoder_position_embeddings (`torch.Tensor` of shape `(batch_size, encoder_seq_len, hidden_size)`): + Position embeddings for keys. **Concatenated** (not added) with key content, doubling the key dimension. + query_position_embeddings (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): + Additional position embeddings. When provided (first layer only), **added** to query content + before concatenation with `query_sine_embed`. Also causes `encoder_position_embeddings` to be + added to key content before concatenation. + attention_mask (`torch.Tensor` of shape `(batch_size, 1, num_queries, encoder_seq_len)`, *optional*): + Attention mask to avoid attending to padding tokens. + """ + query_input_shape = hidden_states.shape[:-1] + kv_input_shape = encoder_hidden_states.shape[:-1] + query_hidden_shape = (*query_input_shape, self.num_attention_heads, self.head_dim) + kv_hidden_shape = (*kv_input_shape, self.num_attention_heads, self.head_dim) + + # Apply content and position projections + query_input = self.q_content_proj(hidden_states) + key_input = self.k_content_proj(encoder_hidden_states) + value_states = self.v_proj(encoder_hidden_states) + key_pos = self.k_pos_proj(encoder_position_embeddings) + + # Combine content and position embeddings + if query_position_embeddings is not None: + query_input = query_input + self.q_pos_proj(query_position_embeddings) + key_input = key_input + key_pos + + # Reshape and concatenate position embeddings (doubling head_dim) + query_input = query_input.view(query_hidden_shape) + key_input = key_input.view(kv_hidden_shape) + query_sine_embed = self.q_pos_sine_proj(query_sine_embed).view(query_hidden_shape) + key_pos = key_pos.view(kv_hidden_shape) + + query_states = torch.cat([query_input, query_sine_embed], dim=-1).view(*query_input_shape, -1) + key_states = torch.cat([key_input, key_pos], dim=-1).view(*kv_input_shape, -1) + + # Reshape for attention computation + expanded_head_dim = query_states.shape[-1] // self.num_attention_heads + query_states = query_states.view(*query_input_shape, self.num_attention_heads, expanded_head_dim).transpose( + 1, 2 + ) + key_states = key_states.view(*kv_input_shape, self.num_attention_heads, expanded_head_dim).transpose(1, 2) + value_states = value_states.view(kv_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(*query_input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class ConditionalDetrMLP(nn.Module): + def __init__(self, config: ConditionalDetrConfig, hidden_size: int, intermediate_size: int): + super().__init__() + self.fc1 = nn.Linear(hidden_size, intermediate_size) + self.fc2 = nn.Linear(intermediate_size, hidden_size) + self.activation_fn = ACT2FN[config.activation_function] + self.activation_dropout = config.activation_dropout + self.dropout = config.dropout + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + 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) + return hidden_states + + +class ConditionalDetrEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: ConditionalDetrConfig): + super().__init__() + self.hidden_size = config.d_model + self.self_attn = ConditionalDetrSelfAttention( + config=config, + hidden_size=self.hidden_size, + num_attention_heads=config.encoder_attention_heads, + dropout=config.attention_dropout, + ) + self.self_attn_layer_norm = nn.LayerNorm(self.hidden_size) + self.dropout = config.dropout + self.mlp = ConditionalDetrMLP(config, self.hidden_size, config.encoder_ffn_dim) + self.final_layer_norm = nn.LayerNorm(self.hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + spatial_position_embeddings: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, hidden_size)` + attention_mask (`torch.FloatTensor`): attention mask of size + `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative + values. + spatial_position_embeddings (`torch.FloatTensor`, *optional*): + Spatial position embeddings (2D positional encodings of image locations), to be added to both + the queries and keys in self-attention (but not to values). + """ + residual = hidden_states + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_embeddings=spatial_position_embeddings, + **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.mlp(hidden_states) + hidden_states = residual + hidden_states + hidden_states = self.final_layer_norm(hidden_states) + + if self.training: + if 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 ConditionalDetrDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: ConditionalDetrConfig): + super().__init__() + self.hidden_size = config.d_model + self.self_attn = ConditionalDetrDecoderSelfAttention( + config=config, + hidden_size=self.hidden_size, + num_attention_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + ) + self.dropout = config.dropout + + self.self_attn_layer_norm = nn.LayerNorm(self.hidden_size) + self.encoder_attn = ConditionalDetrDecoderCrossAttention( + config=config, + hidden_size=self.hidden_size, + num_attention_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + ) + self.encoder_attn_layer_norm = nn.LayerNorm(self.hidden_size) + self.mlp = ConditionalDetrMLP(config, self.hidden_size, config.decoder_ffn_dim) + self.final_layer_norm = nn.LayerNorm(self.hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + spatial_position_embeddings: torch.Tensor | None = None, + query_position_embeddings: torch.Tensor | None = None, + query_sine_embed: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + is_first: bool | None = False, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (`torch.FloatTensor`): attention mask of size + `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative + values. + spatial_position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Spatial position embeddings (2D positional encodings) that are added to the queries and keys in each self-attention layer. + query_position_embeddings (`torch.FloatTensor`, *optional*): + object_queries that are added to the queries and keys + in the self-attention layer. + encoder_hidden_states (`torch.FloatTensor`): + cross attention input to the layer of shape `(seq_len, batch, embed_dim)` + encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size + `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative + values. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + residual = hidden_states + + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + query_position_embeddings=query_position_embeddings, + 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) + + if encoder_hidden_states is not None: + residual = hidden_states + + hidden_states, _ = self.encoder_attn( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + query_sine_embed=query_sine_embed, + encoder_position_embeddings=spatial_position_embeddings, + # Only pass query_position_embeddings for the first layer + query_position_embeddings=query_position_embeddings if is_first else None, + **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.mlp(hidden_states) + hidden_states = residual + hidden_states + hidden_states = self.final_layer_norm(hidden_states) + + return hidden_states + + +class ConditionalDetrMLPPredictionHead(nn.Module): + """ + Very simple multi-layer perceptron (MLP, also called FFN), used to predict the normalized center coordinates, + height and width of a bounding box w.r.t. an image. + + """ + + def __init__(self, input_dim, hidden_dim, output_dim, num_layers): + super().__init__() + self.num_layers = num_layers + h = [hidden_dim] * (num_layers - 1) + self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim])) + + def forward(self, x): + for i, layer in enumerate(self.layers): + x = nn.functional.relu(layer(x)) if i < self.num_layers - 1 else layer(x) + return x + + +class ConditionalDetrConvBlock(nn.Module): + """Basic conv block: Conv3x3 -> GroupNorm -> Activation.""" + + def __init__(self, in_channels: int, out_channels: int, activation: str = "relu"): + super().__init__() + self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) + self.norm = nn.GroupNorm(min(8, out_channels), out_channels) + self.activation = ACT2FN[activation] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.activation(self.norm(self.conv(x))) + + +class ConditionalDetrFPNFusionStage(nn.Module): + """Single FPN fusion stage combining low-resolution features with high-resolution FPN features.""" + + def __init__(self, fpn_channels: int, current_channels: int, output_channels: int, activation: str = "relu"): + super().__init__() + self.fpn_adapter = nn.Conv2d(fpn_channels, current_channels, kernel_size=1) + self.refine = ConditionalDetrConvBlock(current_channels, output_channels, activation) + + def forward(self, features: torch.Tensor, fpn_features: torch.Tensor) -> torch.Tensor: + """ + Args: + features: Current features to upsample, shape (B*Q, current_channels, H_in, W_in) + fpn_features: FPN features at target resolution, shape (B*Q, fpn_channels, H_out, W_out) + + Returns: + Fused and refined features, shape (B*Q, output_channels, H_out, W_out) + """ + fpn_features = self.fpn_adapter(fpn_features) + features = nn.functional.interpolate(features, size=fpn_features.shape[-2:], mode="nearest") + return self.refine(fpn_features + features) + + +class ConditionalDetrMaskHeadSmallConv(nn.Module): + """ + Segmentation mask head that generates per-query masks using FPN-based progressive upsampling. + + Combines attention maps (spatial localization) with encoder features (semantics) and progressively + upsamples through multiple scales, fusing with FPN features for high-resolution detail. + """ + + def __init__( + self, + input_channels: int, + fpn_channels: list[int], + hidden_size: int, + activation_function: str = "relu", + ): + super().__init__() + if input_channels % 8 != 0: + raise ValueError(f"input_channels must be divisible by 8, got {input_channels}") + + self.conv1 = ConditionalDetrConvBlock(input_channels, input_channels, activation_function) + self.conv2 = ConditionalDetrConvBlock(input_channels, hidden_size // 2, activation_function) + + # Progressive channel reduction: /2 -> /4 -> /8 -> /16 + self.fpn_stages = nn.ModuleList( + [ + ConditionalDetrFPNFusionStage( + fpn_channels[0], hidden_size // 2, hidden_size // 4, activation_function + ), + ConditionalDetrFPNFusionStage( + fpn_channels[1], hidden_size // 4, hidden_size // 8, activation_function + ), + ConditionalDetrFPNFusionStage( + fpn_channels[2], hidden_size // 8, hidden_size // 16, activation_function + ), + ] + ) + + self.output_conv = nn.Conv2d(hidden_size // 16, 1, kernel_size=3, padding=1) + + def forward( + self, + features: torch.Tensor, + attention_masks: torch.Tensor, + fpn_features: list[torch.Tensor], + ) -> torch.Tensor: + """ + Args: + features: Encoder output features, shape (batch_size, hidden_size, H, W) + attention_masks: Cross-attention maps from decoder, shape (batch_size, num_queries, num_heads, H, W) + fpn_features: List of 3 FPN features from low to high resolution, each (batch_size, C, H, W) + + Returns: + Predicted masks, shape (batch_size * num_queries, 1, output_H, output_W) + """ + num_queries = attention_masks.shape[1] + + # Expand to (batch_size * num_queries) dimension + features = features.unsqueeze(1).expand(-1, num_queries, -1, -1, -1).flatten(0, 1) + attention_masks = attention_masks.flatten(0, 1) + fpn_features = [ + fpn_feat.unsqueeze(1).expand(-1, num_queries, -1, -1, -1).flatten(0, 1) for fpn_feat in fpn_features + ] + + hidden_states = torch.cat([features, attention_masks], dim=1) + hidden_states = self.conv1(hidden_states) + hidden_states = self.conv2(hidden_states) + + for fpn_stage, fpn_feat in zip(self.fpn_stages, fpn_features): + hidden_states = fpn_stage(hidden_states, fpn_feat) + + return self.output_conv(hidden_states) + + +class ConditionalDetrMHAttentionMap(nn.Module): + """This is a 2D attention module, which only returns the attention softmax (no multiplication by value)""" + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + dropout: float = 0.0, + bias: bool = True, + ): + super().__init__() + self.head_dim = hidden_size // num_attention_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = dropout + + self.q_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + self.k_proj = nn.Linear(hidden_size, hidden_size, bias=bias) + + def forward( + self, query_states: torch.Tensor, key_states: torch.Tensor, attention_mask: torch.Tensor | None = None + ): + query_hidden_shape = (*query_states.shape[:-1], -1, self.head_dim) + key_hidden_shape = (key_states.shape[0], -1, self.head_dim, *key_states.shape[-2:]) + + query_states = self.q_proj(query_states).view(query_hidden_shape) + key_states = nn.functional.conv2d( + key_states, self.k_proj.weight.unsqueeze(-1).unsqueeze(-1), self.k_proj.bias + ).view(key_hidden_shape) + + batch_size, num_queries, num_heads, head_dim = query_states.shape + _, _, _, height, width = key_states.shape + query_shape = (batch_size * num_heads, num_queries, head_dim) + key_shape = (batch_size * num_heads, height * width, head_dim) + attn_weights_shape = (batch_size, num_heads, num_queries, height, width) + + query = query_states.transpose(1, 2).contiguous().view(query_shape) + key = key_states.permute(0, 1, 3, 4, 2).contiguous().view(key_shape) + + attn_weights = ( + (torch.matmul(query * self.scaling, key.transpose(1, 2))).view(attn_weights_shape).transpose(1, 2) + ) + + if attention_mask is not None: + attn_weights = attn_weights + attention_mask + + attn_weights = nn.functional.softmax(attn_weights.flatten(2), dim=-1).view(attn_weights.size()) + attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training) + + return attn_weights + + +@auto_docstring +class ConditionalDetrPreTrainedModel(PreTrainedModel): + config: ConditionalDetrConfig + base_model_prefix = "model" + main_input_name = "pixel_values" + input_modalities = ("image",) + _no_split_modules = [r"ConditionalDetrConvEncoder", r"ConditionalDetrEncoderLayer", r"ConditionalDetrDecoderLayer"] + supports_gradient_checkpointing = True + _supports_sdpa = True + _supports_flash_attn = True + _supports_attention_backend = True + _supports_flex_attn = True # Uses create_bidirectional_masks for attention masking + _keys_to_ignore_on_load_unexpected = [ + r"detr\.model\.backbone\.model\.layer\d+\.0\.downsample\.1\.num_batches_tracked" + ] + + @torch.no_grad() + def _init_weights(self, module): + std = self.config.init_std + xavier_std = self.config.init_xavier_std + + if isinstance(module, ConditionalDetrMaskHeadSmallConv): + # ConditionalDetrMaskHeadSmallConv uses kaiming initialization for all its Conv2d layers + for m in module.modules(): + if isinstance(m, nn.Conv2d): + init.kaiming_uniform_(m.weight, a=1) + if m.bias is not None: + init.constant_(m.bias, 0) + elif isinstance(module, ConditionalDetrMHAttentionMap): + init.zeros_(module.k_proj.bias) + init.zeros_(module.q_proj.bias) + init.xavier_uniform_(module.k_proj.weight, gain=xavier_std) + init.xavier_uniform_(module.q_proj.weight, gain=xavier_std) + elif isinstance(module, ConditionalDetrLearnedPositionEmbedding): + init.uniform_(module.row_embeddings.weight) + init.uniform_(module.column_embeddings.weight) + elif 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]) + elif isinstance(module, (nn.LayerNorm, nn.GroupNorm)): + init.ones_(module.weight) + init.zeros_(module.bias) + + +class ConditionalDetrEncoder(ConditionalDetrPreTrainedModel): + """ + Transformer encoder that processes a flattened feature map from a vision backbone, composed of a stack of + [`ConditionalDetrEncoderLayer`] modules. + + Args: + config (`ConditionalDetrConfig`): Model configuration object. + """ + + _can_record_outputs = {"hidden_states": ConditionalDetrEncoderLayer, "attentions": ConditionalDetrSelfAttention} + + def __init__(self, config: ConditionalDetrConfig): + super().__init__(config) + + self.dropout = config.dropout + self.layers = nn.ModuleList([ConditionalDetrEncoderLayer(config) for _ in range(config.encoder_layers)]) + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + inputs_embeds=None, + attention_mask=None, + spatial_position_embeddings=None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutput: + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Flattened feature map (output of the backbone + projection layer) that is passed to the encoder. + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding pixel features. Mask values selected in `[0, 1]`: + + - 1 for pixel features that are real (i.e. **not masked**), + - 0 for pixel features that are padding (i.e. **masked**). + + [What are attention masks?](../glossary#attention-mask) + spatial_position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Spatial position embeddings (2D positional encodings) that are added to the queries and keys in each self-attention layer. + """ + hidden_states = inputs_embeds + 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 encoder_layer in self.layers: + # we add spatial_position_embeddings as extra input to the encoder_layer + hidden_states = encoder_layer( + hidden_states, attention_mask, spatial_position_embeddings=spatial_position_embeddings, **kwargs + ) + + return BaseModelOutput(last_hidden_state=hidden_states) + + +# function to generate sine positional embedding for 2d coordinates +def gen_sine_position_embeddings(pos_tensor, d_model): + scale = 2 * math.pi + dim = d_model // 2 + dim_t = torch.arange(dim, dtype=torch.float32, device=pos_tensor.device) + dim_t = 10000 ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / dim) + x_embed = pos_tensor[:, :, 0] * scale + y_embed = pos_tensor[:, :, 1] * scale + pos_x = x_embed[:, :, None] / dim_t + pos_y = y_embed[:, :, None] / dim_t + pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2) + pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2) + pos = torch.cat((pos_y, pos_x), dim=2) + return pos.to(pos_tensor.dtype) + + +class ConditionalDetrDecoder(ConditionalDetrPreTrainedModel): + """ + Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`ConditionalDetrDecoderLayer`]. + + The decoder updates the query embeddings through multiple self-attention and cross-attention layers. + + Some small tweaks for Conditional DETR: + + - object_queries and query_position_embeddings are added to the forward pass. + - if self.config.auxiliary_loss is set to True, also returns a stack of activations from all decoding layers. + + Args: + config: ConditionalDetrConfig + """ + + _can_record_outputs = { + "hidden_states": ConditionalDetrDecoderLayer, + "attentions": OutputRecorder(ConditionalDetrDecoderSelfAttention, layer_name="self_attn", index=1), + "cross_attentions": OutputRecorder(ConditionalDetrDecoderCrossAttention, layer_name="encoder_attn", index=1), + } + + def __init__(self, config: ConditionalDetrConfig): + super().__init__(config) + self.hidden_size = config.d_model + + self.dropout = config.dropout + self.layerdrop = config.decoder_layerdrop + + self.layers = nn.ModuleList([ConditionalDetrDecoderLayer(config) for _ in range(config.decoder_layers)]) + # in Conditional DETR, the decoder uses layernorm after the last decoder layer output + self.layernorm = nn.LayerNorm(config.d_model) + + # query_scale is the FFN applied on f to generate transformation T + self.query_scale = ConditionalDetrMLPPredictionHead(self.hidden_size, self.hidden_size, self.hidden_size, 2) + self.ref_point_head = ConditionalDetrMLPPredictionHead(self.hidden_size, self.hidden_size, 2, 2) + for layer_id in range(config.decoder_layers - 1): + # Set q_pos_proj to None for layers after the first (only first layer uses query position embeddings) + self.layers[layer_id + 1].encoder_attn.q_pos_proj = None + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + inputs_embeds=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + spatial_position_embeddings=None, + object_queries_position_embeddings=None, + **kwargs: Unpack[TransformersKwargs], + ) -> ConditionalDetrDecoderOutput: + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + The query embeddings that are passed into the decoder. + + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on certain queries. Mask values selected in `[0, 1]`: + + - 1 for queries that are **not masked**, + - 0 for queries 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 pixel_values of the encoder. Mask values selected + in `[0, 1]`: + + - 1 for pixels that are real (i.e. **not masked**), + - 0 for pixels that are padding (i.e. **masked**). + + spatial_position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Spatial position embeddings that are added to the queries and keys in each cross-attention layer. + object_queries_position_embeddings (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`): + , *optional*): Position embeddings that are added to the queries and keys in each self-attention layer. + """ + if inputs_embeds is not None: + hidden_states = inputs_embeds + + # expand encoder attention mask + if encoder_hidden_states is not None and encoder_attention_mask is not None: + # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len] + encoder_attention_mask = create_bidirectional_mask( + self.config, + inputs_embeds, + encoder_attention_mask, + ) + + # optional intermediate hidden states + intermediate = () if self.config.auxiliary_loss else None + + reference_points_before_sigmoid = self.ref_point_head( + object_queries_position_embeddings + ) # [num_queries, batch_size, 2] + reference_points = reference_points_before_sigmoid.sigmoid().transpose(0, 1) + obj_center = reference_points[..., :2].transpose(0, 1) + # get sine embedding for the query vector + query_sine_embed_before_transformation = gen_sine_position_embeddings(obj_center, self.config.d_model) + + for idx, decoder_layer in enumerate(self.layers): + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: + continue + if idx == 0: + pos_transformation = 1 + else: + pos_transformation = self.query_scale(hidden_states) + # apply transformation + query_sine_embed = query_sine_embed_before_transformation * pos_transformation + + hidden_states = decoder_layer( + hidden_states, + None, + spatial_position_embeddings, + object_queries_position_embeddings, + query_sine_embed, + encoder_hidden_states, # as a positional argument for gradient checkpointing + encoder_attention_mask=encoder_attention_mask, + is_first=(idx == 0), + **kwargs, + ) + + if self.config.auxiliary_loss: + hidden_states = self.layernorm(hidden_states) + intermediate += (hidden_states,) + + # finally, apply layernorm + hidden_states = self.layernorm(hidden_states) + + # stack intermediate decoder activations + if self.config.auxiliary_loss: + intermediate = torch.stack(intermediate) + + return ConditionalDetrDecoderOutput( + last_hidden_state=hidden_states, + intermediate_hidden_states=intermediate, + reference_points=reference_points, + ) + + +@auto_docstring( + custom_intro=""" + The bare CONDITIONAL_DETR Model (consisting of a backbone and encoder-decoder Transformer) outputting raw hidden-states without + any specific head on top. + """ +) +class ConditionalDetrModel(ConditionalDetrPreTrainedModel): + def __init__(self, config: ConditionalDetrConfig): + super().__init__(config) + + self.backbone = ConditionalDetrConvEncoder(config) + + if config.position_embedding_type == "sine": + self.position_embedding = ConditionalDetrSinePositionEmbedding(config.d_model // 2, normalize=True) + elif config.position_embedding_type == "learned": + self.position_embedding = ConditionalDetrLearnedPositionEmbedding(config.d_model // 2) + else: + raise ValueError(f"Not supported {config.position_embedding_type}") + self.query_position_embeddings = nn.Embedding(config.num_queries, config.d_model) + self.input_projection = nn.Conv2d(self.backbone.intermediate_channel_sizes[-1], config.d_model, kernel_size=1) + + self.encoder = ConditionalDetrEncoder(config) + self.decoder = ConditionalDetrDecoder(config) + + # Initialize weights and apply final processing + self.post_init() + + def freeze_backbone(self): + for _, param in self.backbone.model.named_parameters(): + param.requires_grad_(False) + + def unfreeze_backbone(self): + for _, param in self.backbone.model.named_parameters(): + param.requires_grad_(True) + + @auto_docstring + @can_return_tuple + def forward( + self, + pixel_values: torch.FloatTensor, + pixel_mask: torch.LongTensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + encoder_outputs: torch.FloatTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + decoder_inputs_embeds: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> ConditionalDetrModelOutput: + r""" + decoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_queries)`, *optional*): + Not used by default. Can be used to mask object queries. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing the flattened feature map (output of the backbone + projection layer), you + can choose to directly pass a flattened representation of an image. + decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): + Optionally, instead of initializing the queries with a tensor of zeros, you can choose to directly pass an + embedded representation. + + Examples: + + ```python + >>> from transformers import AutoImageProcessor, AutoModel + >>> from PIL import Image + >>> import requests + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/conditional-detr-resnet-50") + >>> model = AutoModel.from_pretrained("microsoft/conditional-detr-resnet-50") + + >>> # prepare image for the model + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> # forward pass + >>> outputs = model(**inputs) + + >>> # the last hidden states are the final query embeddings of the Transformer decoder + >>> # these are of shape (batch_size, num_queries, hidden_size) + >>> last_hidden_states = outputs.last_hidden_state + >>> list(last_hidden_states.shape) + [1, 300, 256] + ```""" + batch_size, num_channels, height, width = pixel_values.shape + device = pixel_values.device + + if pixel_mask is None: + pixel_mask = torch.ones(((batch_size, height, width)), device=device) + + # First, sent pixel_values + pixel_mask through Backbone to obtain the features + # pixel_values should be of shape (batch_size, num_channels, height, width) + # pixel_mask should be of shape (batch_size, height, width) + features = self.backbone(pixel_values, pixel_mask) + + # get final feature map and downsampled mask + feature_map, mask = features[-1] + + if mask is None: + raise ValueError("Backbone does not return downsampled pixel mask") + + # Second, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default) + projected_feature_map = self.input_projection(feature_map) + + # Generate position embeddings + spatial_position_embeddings = self.position_embedding( + shape=feature_map.shape, device=device, dtype=pixel_values.dtype, mask=mask + ) + + # Third, flatten the feature map of shape NxCxHxW to NxCxHW, and permute it to NxHWxC + # In other words, turn their shape into (batch_size, sequence_length, hidden_size) + flattened_features = projected_feature_map.flatten(2).permute(0, 2, 1) + + flattened_mask = mask.flatten(1) + + # Fourth, sent flattened_features + flattened_mask + spatial_position_embeddings through encoder + # flattened_features is a Tensor of shape (batch_size, height*width, hidden_size) + # flattened_mask is a Tensor of shape (batch_size, height*width) + if encoder_outputs is None: + encoder_outputs = self.encoder( + inputs_embeds=flattened_features, + attention_mask=flattened_mask, + spatial_position_embeddings=spatial_position_embeddings, + **kwargs, + ) + # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput + 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, + ) + + # Fifth, sent query embeddings through the decoder (which is conditioned on the encoder output) + object_queries_position_embeddings = self.query_position_embeddings.weight.unsqueeze(0).repeat( + batch_size, 1, 1 + ) + queries = torch.zeros_like(object_queries_position_embeddings) + + # decoder outputs consists of (dec_features, dec_hidden, dec_attn) + decoder_outputs = self.decoder( + inputs_embeds=queries, + attention_mask=None, + spatial_position_embeddings=spatial_position_embeddings, + object_queries_position_embeddings=object_queries_position_embeddings, + encoder_hidden_states=encoder_outputs.last_hidden_state, + encoder_attention_mask=flattened_mask, + **kwargs, + ) + + return ConditionalDetrModelOutput( + last_hidden_state=decoder_outputs.last_hidden_state, + 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, + intermediate_hidden_states=decoder_outputs.intermediate_hidden_states, + reference_points=decoder_outputs.reference_points, + ) + + +def inverse_sigmoid(x, eps=1e-5): + x = x.clamp(min=0, max=1) + x1 = x.clamp(min=eps) + x2 = (1 - x).clamp(min=eps) + return torch.log(x1 / x2) + + +@auto_docstring( + custom_intro=""" + CONDITIONAL_DETR Model (consisting of a backbone and encoder-decoder Transformer) with object detection heads on top, for tasks + such as COCO detection. + """ +) +class ConditionalDetrForObjectDetection(ConditionalDetrPreTrainedModel): + def __init__(self, config: ConditionalDetrConfig): + super().__init__(config) + + # CONDITIONAL_DETR encoder-decoder model + self.model = ConditionalDetrModel(config) + self.class_labels_classifier = nn.Linear(config.d_model, config.num_labels) + self.bbox_predictor = ConditionalDetrMLPPredictionHead( + input_dim=config.d_model, hidden_dim=config.d_model, output_dim=4, num_layers=3 + ) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + @can_return_tuple + def forward( + self, + pixel_values: torch.FloatTensor, + pixel_mask: torch.LongTensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + encoder_outputs: torch.FloatTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + decoder_inputs_embeds: torch.FloatTensor | None = None, + labels: list[dict] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> ConditionalDetrObjectDetectionOutput: + r""" + decoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_queries)`, *optional*): + Not used by default. Can be used to mask object queries. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing the flattened feature map (output of the backbone + projection layer), you + can choose to directly pass a flattened representation of an image. + decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): + Optionally, instead of initializing the queries with a tensor of zeros, you can choose to directly pass an + embedded representation. + labels (`list[Dict]` of len `(batch_size,)`, *optional*): + Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the + following 2 keys: 'class_labels' and 'boxes' (the class labels and bounding boxes of an image in the batch + respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes + in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)`. + + Examples: + + ```python + >>> from transformers import AutoImageProcessor, AutoModelForObjectDetection + >>> from PIL import Image + >>> import requests + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/conditional-detr-resnet-50") + >>> model = AutoModelForObjectDetection.from_pretrained("microsoft/conditional-detr-resnet-50") + + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> outputs = model(**inputs) + + >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax) + >>> target_sizes = torch.tensor([image.size[::-1]]) + >>> results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[ + ... 0 + ... ] + >>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): + ... box = [round(i, 2) for i in box.tolist()] + ... print( + ... f"Detected {model.config.id2label[label.item()]} with confidence " + ... f"{round(score.item(), 3)} at location {box}" + ... ) + Detected remote with confidence 0.833 at location [38.31, 72.1, 177.63, 118.45] + Detected cat with confidence 0.831 at location [9.2, 51.38, 321.13, 469.0] + Detected cat with confidence 0.804 at location [340.3, 16.85, 642.93, 370.95] + Detected remote with confidence 0.683 at location [334.48, 73.49, 366.37, 190.01] + Detected couch with confidence 0.535 at location [0.52, 1.19, 640.35, 475.1] + ```""" + # First, sent images through CONDITIONAL_DETR base model to obtain encoder + decoder outputs + outputs = self.model( + pixel_values, + pixel_mask=pixel_mask, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + **kwargs, + ) + + sequence_output = outputs[0] + + # class logits + predicted bounding boxes + logits = self.class_labels_classifier(sequence_output) + + reference = outputs.reference_points + reference_before_sigmoid = inverse_sigmoid(reference).transpose(0, 1) + + hs = sequence_output + tmp = self.bbox_predictor(hs) + tmp[..., :2] += reference_before_sigmoid + pred_boxes = tmp.sigmoid() + # pred_boxes = self.bbox_predictor(sequence_output).sigmoid() + + loss, loss_dict, auxiliary_outputs = None, None, None + if labels is not None: + outputs_class, outputs_coord = None, None + if self.config.auxiliary_loss: + outputs_coords = [] + intermediate = outputs.intermediate_hidden_states + outputs_class = self.class_labels_classifier(intermediate) + for lvl in range(intermediate.shape[0]): + tmp = self.bbox_predictor(intermediate[lvl]) + tmp[..., :2] += reference_before_sigmoid + outputs_coord = tmp.sigmoid() + outputs_coords.append(outputs_coord) + outputs_coord = torch.stack(outputs_coords) + loss, loss_dict, auxiliary_outputs = self.loss_function( + logits, labels, self.device, pred_boxes, self.config, outputs_class, outputs_coord + ) + + return ConditionalDetrObjectDetectionOutput( + loss=loss, + loss_dict=loss_dict, + logits=logits, + pred_boxes=pred_boxes, + auxiliary_outputs=auxiliary_outputs, + last_hidden_state=outputs.last_hidden_state, + 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, + ) + + # taken from https://github.com/Atten4Vis/conditionalDETR/blob/master/models/conditional_detr.py + def _set_aux_loss(self, outputs_class, outputs_coord): + return [{"logits": a, "pred_boxes": b} for a, b in zip(outputs_class[:-1], outputs_coord[:-1])] + + +@auto_docstring( + custom_intro=""" + CONDITIONAL_DETR Model (consisting of a backbone and encoder-decoder Transformer) with a segmentation head on top, for tasks + such as COCO panoptic. + """ +) +class ConditionalDetrForSegmentation(ConditionalDetrPreTrainedModel): + def __init__(self, config: ConditionalDetrConfig): + super().__init__(config) + + # object detection model + self.conditional_detr = ConditionalDetrForObjectDetection(config) + + # segmentation head + hidden_size, number_of_heads = config.d_model, config.encoder_attention_heads + intermediate_channel_sizes = self.conditional_detr.model.backbone.intermediate_channel_sizes + + self.mask_head = ConditionalDetrMaskHeadSmallConv( + input_channels=hidden_size + number_of_heads, + fpn_channels=intermediate_channel_sizes[::-1][-3:], + hidden_size=hidden_size, + activation_function=config.activation_function, + ) + + self.bbox_attention = ConditionalDetrMHAttentionMap(hidden_size, number_of_heads, dropout=0.0) + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + @can_return_tuple + def forward( + self, + pixel_values: torch.FloatTensor, + pixel_mask: torch.LongTensor | None = None, + decoder_attention_mask: torch.FloatTensor | None = None, + encoder_outputs: torch.FloatTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + decoder_inputs_embeds: torch.FloatTensor | None = None, + labels: list[dict] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.FloatTensor] | ConditionalDetrSegmentationOutput: + r""" + decoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_queries)`, *optional*): + Mask to avoid performing attention on certain object queries in the decoder. Mask values selected in `[0, 1]`: + + - 1 for queries that are **not masked**, + - 0 for queries that are **masked**. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Kept for backward compatibility, but cannot be used for segmentation, as segmentation requires + multi-scale features from the backbone that are not available when bypassing it with inputs_embeds. + decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): + Optionally, instead of initializing the queries with a tensor of zeros, you can choose to directly pass an + embedded representation. Useful for tasks that require custom query initialization. + labels (`list[Dict]` of len `(batch_size,)`, *optional*): + Labels for computing the bipartite matching loss, DICE/F-1 loss and Focal loss. List of dicts, each + dictionary containing at least the following 3 keys: 'class_labels', 'boxes' and 'masks' (the class labels, + bounding boxes and segmentation masks of an image in the batch respectively). The class labels themselves + should be a `torch.LongTensor` of len `(number of bounding boxes in the image,)`, the boxes a + `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)` and the masks a + `torch.FloatTensor` of shape `(number of bounding boxes in the image, height, width)`. + + Examples: + + ```python + >>> import io + >>> import httpx + >>> from io import BytesIO + >>> from PIL import Image + >>> import torch + >>> import numpy + + >>> from transformers import AutoImageProcessor, ConditionalDetrForSegmentation + >>> from transformers.image_transforms import rgb_to_id + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> with httpx.stream("GET", url) as response: + ... image = Image.open(BytesIO(response.read())) + + >>> image_processor = AutoImageProcessor.from_pretrained("facebook/conditional_detr-resnet-50-panoptic") + >>> model = ConditionalDetrForSegmentation.from_pretrained("facebook/conditional_detr-resnet-50-panoptic") + + >>> # prepare image for the model + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> # forward pass + >>> outputs = model(**inputs) + + >>> # Use the `post_process_panoptic_segmentation` method of the `image_processor` to retrieve post-processed panoptic segmentation maps + >>> # Segmentation results are returned as a list of dictionaries + >>> result = image_processor.post_process_panoptic_segmentation(outputs, target_sizes=[(300, 500)]) + + >>> # A tensor of shape (height, width) where each value denotes a segment id, filled with -1 if no segment is found + >>> panoptic_seg = result[0]["segmentation"] + >>> panoptic_seg.shape + torch.Size([300, 500]) + >>> # Get prediction score and segment_id to class_id mapping of each segment + >>> panoptic_segments_info = result[0]["segments_info"] + >>> len(panoptic_segments_info) + 5 + ```""" + + batch_size, num_channels, height, width = pixel_values.shape + device = pixel_values.device + + if pixel_mask is None: + pixel_mask = torch.ones((batch_size, height, width), device=device) + + vision_features = self.conditional_detr.model.backbone(pixel_values, pixel_mask) + feature_map, mask = vision_features[-1] + + # Apply 1x1 conv to map (batch_size, C, H, W) -> (batch_size, hidden_size, H, W), then flatten to (batch_size, HW, hidden_size) + projected_feature_map = self.conditional_detr.model.input_projection(feature_map) + flattened_features = projected_feature_map.flatten(2).permute(0, 2, 1) + spatial_position_embeddings = self.conditional_detr.model.position_embedding( + shape=feature_map.shape, device=device, dtype=pixel_values.dtype, mask=mask + ) + flattened_mask = mask.flatten(1) + + if encoder_outputs is None: + encoder_outputs = self.conditional_detr.model.encoder( + inputs_embeds=flattened_features, + attention_mask=flattened_mask, + spatial_position_embeddings=spatial_position_embeddings, + **kwargs, + ) + + object_queries_position_embeddings = self.conditional_detr.model.query_position_embeddings.weight.unsqueeze( + 0 + ).repeat(batch_size, 1, 1) + + # Use decoder_inputs_embeds as queries if provided, otherwise initialize with zeros + if decoder_inputs_embeds is not None: + queries = decoder_inputs_embeds + else: + queries = torch.zeros_like(object_queries_position_embeddings) + + decoder_outputs = self.conditional_detr.model.decoder( + inputs_embeds=queries, + attention_mask=decoder_attention_mask, + spatial_position_embeddings=spatial_position_embeddings, + object_queries_position_embeddings=object_queries_position_embeddings, + encoder_hidden_states=encoder_outputs.last_hidden_state, + encoder_attention_mask=flattened_mask, + **kwargs, + ) + + sequence_output = decoder_outputs[0] + + logits = self.conditional_detr.class_labels_classifier(sequence_output) + pred_boxes = self.conditional_detr.bbox_predictor(sequence_output).sigmoid() + + height, width = feature_map.shape[-2:] + memory = encoder_outputs.last_hidden_state.permute(0, 2, 1).view( + batch_size, self.config.d_model, height, width + ) + attention_mask = flattened_mask.view(batch_size, height, width) + + if attention_mask is not None: + min_dtype = torch.finfo(memory.dtype).min + attention_mask = torch.where( + attention_mask.unsqueeze(1).unsqueeze(1), + torch.tensor(0.0, device=memory.device, dtype=memory.dtype), + min_dtype, + ) + + bbox_mask = self.bbox_attention(sequence_output, memory, attention_mask=attention_mask) + + seg_masks = self.mask_head( + features=projected_feature_map, + attention_masks=bbox_mask, + fpn_features=[vision_features[2][0], vision_features[1][0], vision_features[0][0]], + ) + + pred_masks = seg_masks.view( + batch_size, self.conditional_detr.config.num_queries, seg_masks.shape[-2], seg_masks.shape[-1] + ) + + loss, loss_dict, auxiliary_outputs = None, None, None + if labels is not None: + outputs_class, outputs_coord = None, None + if self.config.auxiliary_loss: + intermediate = decoder_outputs.intermediate_hidden_states + outputs_class = self.conditional_detr.class_labels_classifier(intermediate) + outputs_coord = self.conditional_detr.bbox_predictor(intermediate).sigmoid() + loss, loss_dict, auxiliary_outputs = self.loss_function( + logits, labels, device, pred_boxes, pred_masks, self.config, outputs_class, outputs_coord + ) + + return ConditionalDetrSegmentationOutput( + loss=loss, + loss_dict=loss_dict, + logits=logits, + pred_boxes=pred_boxes, + pred_masks=pred_masks, + auxiliary_outputs=auxiliary_outputs, + last_hidden_state=decoder_outputs.last_hidden_state, + 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, + ) + + +__all__ = [ + "ConditionalDetrForObjectDetection", + "ConditionalDetrForSegmentation", + "ConditionalDetrModel", + "ConditionalDetrPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/conditional_detr/modular_conditional_detr.py b/third_party/transformers/src/transformers/models/conditional_detr/modular_conditional_detr.py new file mode 100644 index 0000000000000000000000000000000000000000..ffc1e78bee01ce583a3fc8abca32e87cb7798814 --- /dev/null +++ b/third_party/transformers/src/transformers/models/conditional_detr/modular_conditional_detr.py @@ -0,0 +1,1109 @@ +# Copyright 2022 Microsoft Research Asia 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 math +from collections.abc import Callable + +import torch +from torch import nn + +from ...image_transforms import ( + center_to_corners_format, +) +from ...image_utils import AnnotationFormat +from ...masking_utils import create_bidirectional_mask +from ...modeling_outputs import ( + BaseModelOutput, +) +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import ( + TensorType, + TransformersKwargs, + auto_docstring, + logging, + requires_backends, +) +from ...utils.generic import can_return_tuple, merge_with_config_defaults +from ...utils.import_utils import requires +from ...utils.output_capturing import OutputRecorder, capture_outputs +from ..deformable_detr.modeling_deformable_detr import inverse_sigmoid +from ..detr.image_processing_detr import DetrImageProcessor +from ..detr.image_processing_pil_detr import DetrImageProcessorPil +from ..detr.modeling_detr import ( + DetrConvEncoder, + DetrDecoderLayer, + DetrDecoderOutput, + DetrEncoder, + DetrEncoderLayer, + DetrForObjectDetection, + DetrForSegmentation, + DetrLearnedPositionEmbedding, + DetrMLP, + DetrMLPPredictionHead, + DetrModel, + DetrModelOutput, + DetrObjectDetectionOutput, + DetrPreTrainedModel, + DetrSegmentationOutput, + DetrSelfAttention, + DetrSinePositionEmbedding, + eager_attention_forward, +) +from .configuration_conditional_detr import ConditionalDetrConfig + + +logger = logging.get_logger(__name__) + + +class ConditionalDetrImageProcessorKwargs(ImagesKwargs, total=False): + r""" + format (`str`, *optional*, defaults to `AnnotationFormat.COCO_DETECTION`): + Data format of the annotations. One of "coco_detection" or "coco_panoptic". + do_convert_annotations (`bool`, *optional*, defaults to `True`): + Controls whether to convert the annotations to the format expected by the CONDITIONAL_DETR model. Converts the + bounding boxes to the format `(center_x, center_y, width, height)` and in the range `[0, 1]`. + Can be overridden by the `do_convert_annotations` parameter in the `preprocess` method. + """ + + format: str | AnnotationFormat + do_convert_annotations: bool + + +class ConditionalDetrImageProcessor(DetrImageProcessor): + def post_process_object_detection( + self, outputs, threshold: float = 0.5, target_sizes: TensorType | list[tuple] = None, top_k: int = 100 + ): + """ + Converts the raw output of [`ConditionalDetrForObjectDetection`] into final bounding boxes in (top_left_x, + top_left_y, bottom_right_x, bottom_right_y) format. Only supports PyTorch. + + Args: + outputs ([`ConditionalDetrObjectDetectionOutput`]): + Raw outputs of the model. + threshold (`float`, *optional*): + Score threshold to keep object detection predictions. + target_sizes (`torch.Tensor` or `list[tuple[int, int]]`, *optional*): + Tensor of shape `(batch_size, 2)` or list of tuples (`tuple[int, int]`) containing the target size + (height, width) of each image in the batch. If left to None, predictions will not be resized. + top_k (`int`, *optional*, defaults to 100): + Keep only top k bounding boxes before filtering by thresholding. + + Returns: + `list[Dict]`: A list of dictionaries, each dictionary containing the scores, labels and boxes for an image + in the batch as predicted by the model. + """ + out_logits, out_bbox = outputs.logits, outputs.pred_boxes + + if target_sizes is not None: + if len(out_logits) != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + prob = out_logits.sigmoid() + prob = prob.view(out_logits.shape[0], -1) + k_value = min(top_k, prob.size(1)) + topk_values, topk_indexes = torch.topk(prob, k_value, dim=1) + scores = topk_values + topk_boxes = torch.div(topk_indexes, out_logits.shape[2], rounding_mode="floor") + labels = topk_indexes % out_logits.shape[2] + boxes = center_to_corners_format(out_bbox) + boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4)) + + # and from relative [0, 1] to absolute [0, height] coordinates + if target_sizes is not None: + if isinstance(target_sizes, list): + img_h = torch.Tensor([i[0] for i in target_sizes]) + img_w = torch.Tensor([i[1] for i in target_sizes]) + else: + img_h, img_w = target_sizes.unbind(1) + scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(boxes.device) + boxes = boxes * scale_fct[:, None, :] + + results = [] + for s, l, b in zip(scores, labels, boxes): + score = s[s > threshold] + label = l[s > threshold] + box = b[s > threshold] + results.append({"scores": score, "labels": label, "boxes": box}) + + return results + + def post_process_semantic_segmentation(self, outputs, target_sizes: list[tuple[int, int]] | None = None): + """ + Converts the output of [`ConditionalDetrForSegmentation`] into semantic segmentation maps. Only supports PyTorch. + + Args: + outputs ([`ConditionalDetrForSegmentation`]): + Raw outputs of the model. + target_sizes (`list[tuple[int, int]]`, *optional*): + A list of tuples (`tuple[int, int]`) containing the target size (height, width) of each image in the + batch. If unset, predictions will not be resized. + Returns: + `list[torch.Tensor]`: + A list of length `batch_size`, where each item is a semantic segmentation map of shape (height, width) + corresponding to the target_sizes entry (if `target_sizes` is specified). Each entry of each + `torch.Tensor` correspond to a semantic class id. + """ + class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes] + masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width] + + # Conditional DETR does not have a null class, so we use all classes + masks_classes = class_queries_logits.softmax(dim=-1) + masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # Semantic segmentation logits of shape (batch_size, num_classes, height, width) + segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs) + batch_size = class_queries_logits.shape[0] + + # Resize logits and compute semantic segmentation maps + if target_sizes is not None: + if batch_size != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + semantic_segmentation = [] + for idx in range(batch_size): + resized_logits = nn.functional.interpolate( + segmentation[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False + ) + semantic_map = resized_logits[0].argmax(dim=0) + semantic_segmentation.append(semantic_map) + else: + semantic_segmentation = segmentation.argmax(dim=1) + semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])] + + return semantic_segmentation + + +class ConditionalDetrImageProcessorPil(DetrImageProcessorPil): + @requires(backends=("torch",)) + def post_process_object_detection( + self, outputs, threshold: float = 0.5, target_sizes: TensorType | list[tuple] = None, top_k: int = 100 + ): + """ + Converts the raw output of [`ConditionalDetrForObjectDetection`] into final bounding boxes in (top_left_x, + top_left_y, bottom_right_x, bottom_right_y) format. Only supports PyTorch. + + Args: + outputs ([`ConditionalDetrObjectDetectionOutput`]): + Raw outputs of the model. + threshold (`float`, *optional*): + Score threshold to keep object detection predictions. + target_sizes (`torch.Tensor` or `list[tuple[int, int]]`, *optional*): + Tensor of shape `(batch_size, 2)` or list of tuples (`tuple[int, int]`) containing the target size + (height, width) of each image in the batch. If left to None, predictions will not be resized. + top_k (`int`, *optional*, defaults to 100): + Keep only top k bounding boxes before filtering by thresholding. + + Returns: + `list[Dict]`: A list of dictionaries, each dictionary containing the scores, labels and boxes for an image + in the batch as predicted by the model. + """ + requires_backends(self, ["torch"]) + out_logits, out_bbox = outputs.logits, outputs.pred_boxes + + if target_sizes is not None: + if len(out_logits) != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + prob = out_logits.sigmoid() + prob = prob.view(out_logits.shape[0], -1) + k_value = min(top_k, prob.size(1)) + topk_values, topk_indexes = torch.topk(prob, k_value, dim=1) + scores = topk_values + topk_boxes = torch.div(topk_indexes, out_logits.shape[2], rounding_mode="floor") + labels = topk_indexes % out_logits.shape[2] + boxes = center_to_corners_format(out_bbox) + boxes = torch.gather(boxes, 1, topk_boxes.unsqueeze(-1).repeat(1, 1, 4)) + + # and from relative [0, 1] to absolute [0, height] coordinates + if target_sizes is not None: + if isinstance(target_sizes, list): + img_h = torch.Tensor([i[0] for i in target_sizes]) + img_w = torch.Tensor([i[1] for i in target_sizes]) + else: + img_h, img_w = target_sizes.unbind(1) + scale_fct = torch.stack([img_w, img_h, img_w, img_h], dim=1).to(boxes.device) + boxes = boxes * scale_fct[:, None, :] + + results = [] + for s, l, b in zip(scores, labels, boxes): + score = s[s > threshold] + label = l[s > threshold] + box = b[s > threshold] + results.append({"scores": score, "labels": label, "boxes": box}) + + return results + + @requires(backends=("torch",)) + def post_process_semantic_segmentation(self, outputs, target_sizes: list[tuple[int, int]] | None = None): + """ + Converts the output of [`ConditionalDetrForSegmentation`] into semantic segmentation maps. Only supports PyTorch. + + Args: + outputs ([`ConditionalDetrForSegmentation`]): + Raw outputs of the model. + target_sizes (`list[tuple[int, int]]`, *optional*): + A list of tuples (`tuple[int, int]`) containing the target size (height, width) of each image in the + batch. If unset, predictions will not be resized. + Returns: + `list[torch.Tensor]`: + A list of length `batch_size`, where each item is a semantic segmentation map of shape (height, width) + corresponding to the target_sizes entry (if `target_sizes` is specified). Each entry of each + `torch.Tensor` correspond to a semantic class id. + """ + requires_backends(self, ["torch"]) + class_queries_logits = outputs.logits # [batch_size, num_queries, num_classes] + masks_queries_logits = outputs.pred_masks # [batch_size, num_queries, height, width] + + # Conditional DETR does not have a null class, so we use all classes + masks_classes = class_queries_logits.softmax(dim=-1) + masks_probs = masks_queries_logits.sigmoid() # [batch_size, num_queries, height, width] + + # Semantic segmentation logits of shape (batch_size, num_classes, height, width) + segmentation = torch.einsum("bqc, bqhw -> bchw", masks_classes, masks_probs) + batch_size = class_queries_logits.shape[0] + + # Resize logits and compute semantic segmentation maps + if target_sizes is not None: + if batch_size != len(target_sizes): + raise ValueError( + "Make sure that you pass in as many target sizes as the batch dimension of the logits" + ) + + semantic_segmentation = [] + for idx in range(batch_size): + resized_logits = nn.functional.interpolate( + segmentation[idx].unsqueeze(dim=0), size=target_sizes[idx], mode="bilinear", align_corners=False + ) + semantic_map = resized_logits[0].argmax(dim=0) + semantic_segmentation.append(semantic_map) + else: + semantic_segmentation = segmentation.argmax(dim=1) + semantic_segmentation = [semantic_segmentation[i] for i in range(semantic_segmentation.shape[0])] + + return semantic_segmentation + + +class ConditionalDetrDecoderOutput(DetrDecoderOutput): + r""" + 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. + intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`): + Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a + layernorm. + reference_points (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, 2 (anchor points))`): + Reference points (reference points of each layer of the decoder). + """ + + reference_points: tuple[torch.FloatTensor] | None = None + + +class ConditionalDetrModelOutput(DetrModelOutput): + 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. + intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, sequence_length, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`): + Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a + layernorm. + reference_points (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, 2 (anchor points))`): + Reference points (reference points of each layer of the decoder). + """ + + reference_points: tuple[torch.FloatTensor] | None = None + + +# function to generate sine positional embedding for 2d coordinates +def gen_sine_position_embeddings(pos_tensor, d_model): + scale = 2 * math.pi + dim = d_model // 2 + dim_t = torch.arange(dim, dtype=torch.float32, device=pos_tensor.device) + dim_t = 10000 ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / dim) + x_embed = pos_tensor[:, :, 0] * scale + y_embed = pos_tensor[:, :, 1] * scale + pos_x = x_embed[:, :, None] / dim_t + pos_y = y_embed[:, :, None] / dim_t + pos_x = torch.stack((pos_x[:, :, 0::2].sin(), pos_x[:, :, 1::2].cos()), dim=3).flatten(2) + pos_y = torch.stack((pos_y[:, :, 0::2].sin(), pos_y[:, :, 1::2].cos()), dim=3).flatten(2) + pos = torch.cat((pos_y, pos_x), dim=2) + return pos.to(pos_tensor.dtype) + + +class ConditionalDetrObjectDetectionOutput(DetrObjectDetectionOutput): + pass + + +class ConditionalDetrSegmentationOutput(DetrSegmentationOutput): + pass + + +class ConditionalDetrConvEncoder(DetrConvEncoder): + pass + + +class ConditionalDetrSinePositionEmbedding(DetrSinePositionEmbedding): + pass + + +class ConditionalDetrLearnedPositionEmbedding(DetrLearnedPositionEmbedding): + pass + + +class ConditionalDetrSelfAttention(DetrSelfAttention): + pass + + +class ConditionalDetrDecoderSelfAttention(nn.Module): + """ + Multi-headed self-attention for Conditional DETR decoder layers. + + This attention module handles separate content and position projections, which are then combined + before applying standard self-attention. Position embeddings are added to both queries and keys. + """ + + def __init__( + self, + config: ConditionalDetrConfig, + hidden_size: int, + num_attention_heads: int, + dropout: float | int = 0.0, + ): + super().__init__() + self.config = config + self.hidden_size = hidden_size + self.head_dim = hidden_size // num_attention_heads + self.scaling = self.head_dim**-0.5 + self.attention_dropout = dropout + self.is_causal = False + + # Content and position projections + self.q_content_proj = nn.Linear(hidden_size, hidden_size) + self.q_pos_proj = nn.Linear(hidden_size, hidden_size) + self.k_content_proj = nn.Linear(hidden_size, hidden_size) + self.k_pos_proj = nn.Linear(hidden_size, hidden_size) + self.v_proj = nn.Linear(hidden_size, hidden_size) + self.o_proj = nn.Linear(hidden_size, hidden_size) + + def forward( + self, + hidden_states: torch.Tensor, + query_position_embeddings: torch.Tensor, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`): + Input hidden states from the decoder layer. + query_position_embeddings (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`): + Position embeddings for queries and keys. Required (unlike standard attention). Processed through + separate position projections (`q_pos_proj`, `k_pos_proj`) and added to content projections. + attention_mask (`torch.Tensor` of shape `(batch_size, 1, num_queries, num_queries)`, *optional*): + Attention mask to avoid attending to padding tokens. + """ + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = ( + (self.q_content_proj(hidden_states) + self.q_pos_proj(query_position_embeddings)) + .view(hidden_shape) + .transpose(1, 2) + ) + key_states = ( + (self.k_content_proj(hidden_states) + self.k_pos_proj(query_position_embeddings)) + .view(hidden_shape) + .transpose(1, 2) + ) + value_states = self.v_proj(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() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class ConditionalDetrDecoderCrossAttention(nn.Module): + """ + Multi-headed cross-attention for Conditional DETR decoder layers. + + This attention module handles the special cross-attention logic in Conditional DETR: + - Separate content and position projections for queries and keys + - Concatenation of query sine embeddings with queries (doubling query dimension) + - Concatenation of key position embeddings with keys (doubling key dimension) + - Output dimension remains hidden_size despite doubled input dimensions + """ + + def __init__( + self, + config: ConditionalDetrConfig, + hidden_size: int, + num_attention_heads: int, + dropout: float | int = 0.0, + ): + super().__init__() + self.config = config + self.hidden_size = hidden_size + self.num_attention_heads = num_attention_heads + self.head_dim = hidden_size // num_attention_heads + self.attention_dropout = dropout + self.is_causal = False + + # Content and position projections + self.q_content_proj = nn.Linear(hidden_size, hidden_size) + self.q_pos_proj = nn.Linear(hidden_size, hidden_size) + self.k_content_proj = nn.Linear(hidden_size, hidden_size) + self.k_pos_proj = nn.Linear(hidden_size, hidden_size) + self.v_proj = nn.Linear(hidden_size, hidden_size) + self.q_pos_sine_proj = nn.Linear(hidden_size, hidden_size) + + # Output projection: input is hidden_size * 2 (from concatenated q/k), output is hidden_size + self.o_proj = nn.Linear(hidden_size, hidden_size) + + # Compute scaling for expanded head_dim (q and k have doubled dimensions after concatenation) + # This matches the original Conditional DETR implementation where embed_dim * 2 is used + expanded_head_dim = (hidden_size * 2) // num_attention_heads + self.scaling = expanded_head_dim**-0.5 + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + query_sine_embed: torch.Tensor, + encoder_position_embeddings: torch.Tensor, + query_position_embeddings: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Args: + hidden_states (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`): + Decoder hidden states (queries). + encoder_hidden_states (`torch.Tensor` of shape `(batch_size, encoder_seq_len, hidden_size)`): + Encoder output hidden states (keys and values). + query_sine_embed (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`): + Sine position embeddings for queries. **Concatenated** (not added) with query content, + doubling the query dimension. + encoder_position_embeddings (`torch.Tensor` of shape `(batch_size, encoder_seq_len, hidden_size)`): + Position embeddings for keys. **Concatenated** (not added) with key content, doubling the key dimension. + query_position_embeddings (`torch.Tensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): + Additional position embeddings. When provided (first layer only), **added** to query content + before concatenation with `query_sine_embed`. Also causes `encoder_position_embeddings` to be + added to key content before concatenation. + attention_mask (`torch.Tensor` of shape `(batch_size, 1, num_queries, encoder_seq_len)`, *optional*): + Attention mask to avoid attending to padding tokens. + """ + query_input_shape = hidden_states.shape[:-1] + kv_input_shape = encoder_hidden_states.shape[:-1] + query_hidden_shape = (*query_input_shape, self.num_attention_heads, self.head_dim) + kv_hidden_shape = (*kv_input_shape, self.num_attention_heads, self.head_dim) + + # Apply content and position projections + query_input = self.q_content_proj(hidden_states) + key_input = self.k_content_proj(encoder_hidden_states) + value_states = self.v_proj(encoder_hidden_states) + key_pos = self.k_pos_proj(encoder_position_embeddings) + + # Combine content and position embeddings + if query_position_embeddings is not None: + query_input = query_input + self.q_pos_proj(query_position_embeddings) + key_input = key_input + key_pos + + # Reshape and concatenate position embeddings (doubling head_dim) + query_input = query_input.view(query_hidden_shape) + key_input = key_input.view(kv_hidden_shape) + query_sine_embed = self.q_pos_sine_proj(query_sine_embed).view(query_hidden_shape) + key_pos = key_pos.view(kv_hidden_shape) + + query_states = torch.cat([query_input, query_sine_embed], dim=-1).view(*query_input_shape, -1) + key_states = torch.cat([key_input, key_pos], dim=-1).view(*kv_input_shape, -1) + + # Reshape for attention computation + expanded_head_dim = query_states.shape[-1] // self.num_attention_heads + query_states = query_states.view(*query_input_shape, self.num_attention_heads, expanded_head_dim).transpose( + 1, 2 + ) + key_states = key_states.view(*kv_input_shape, self.num_attention_heads, expanded_head_dim).transpose(1, 2) + value_states = value_states.view(kv_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(*query_input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class ConditionalDetrMLP(DetrMLP): + pass + + +class ConditionalDetrEncoderLayer(DetrEncoderLayer): + pass + + +class ConditionalDetrDecoderLayer(DetrDecoderLayer): + def __init__(self, config: ConditionalDetrConfig): + super().__init__() + self.self_attn = ConditionalDetrDecoderSelfAttention( + config=config, + hidden_size=self.hidden_size, + num_attention_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + ) + self.encoder_attn = ConditionalDetrDecoderCrossAttention( + config=config, + hidden_size=self.hidden_size, + num_attention_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + spatial_position_embeddings: torch.Tensor | None = None, + query_position_embeddings: torch.Tensor | None = None, + query_sine_embed: torch.Tensor | None = None, + encoder_hidden_states: torch.Tensor | None = None, + encoder_attention_mask: torch.Tensor | None = None, + is_first: bool | None = False, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + """ + Args: + hidden_states (`torch.FloatTensor`): input to the layer of shape `(seq_len, batch, embed_dim)` + attention_mask (`torch.FloatTensor`): attention mask of size + `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative + values. + spatial_position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + Spatial position embeddings (2D positional encodings) that are added to the queries and keys in each self-attention layer. + query_position_embeddings (`torch.FloatTensor`, *optional*): + object_queries that are added to the queries and keys + in the self-attention layer. + encoder_hidden_states (`torch.FloatTensor`): + cross attention input to the layer of shape `(seq_len, batch, embed_dim)` + encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size + `(batch, 1, target_len, source_len)` where padding elements are indicated by very large negative + values. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + """ + residual = hidden_states + + hidden_states, _ = self.self_attn( + hidden_states=hidden_states, + query_position_embeddings=query_position_embeddings, + 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) + + if encoder_hidden_states is not None: + residual = hidden_states + + hidden_states, _ = self.encoder_attn( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + query_sine_embed=query_sine_embed, + encoder_position_embeddings=spatial_position_embeddings, + # Only pass query_position_embeddings for the first layer + query_position_embeddings=query_position_embeddings if is_first else None, + **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.mlp(hidden_states) + hidden_states = residual + hidden_states + hidden_states = self.final_layer_norm(hidden_states) + + return hidden_states + + +class ConditionalDetrMLPPredictionHead(DetrMLPPredictionHead): + pass + + +class ConditionalDetrPreTrainedModel(DetrPreTrainedModel): + _keys_to_ignore_on_load_unexpected = [ + r"detr\.model\.backbone\.model\.layer\d+\.0\.downsample\.1\.num_batches_tracked" + ] + + +class ConditionalDetrEncoder(DetrEncoder): + pass + + +class ConditionalDetrDecoder(ConditionalDetrPreTrainedModel): + """ + Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`ConditionalDetrDecoderLayer`]. + + The decoder updates the query embeddings through multiple self-attention and cross-attention layers. + + Some small tweaks for Conditional DETR: + + - object_queries and query_position_embeddings are added to the forward pass. + - if self.config.auxiliary_loss is set to True, also returns a stack of activations from all decoding layers. + + Args: + config: ConditionalDetrConfig + """ + + _can_record_outputs = { + "hidden_states": ConditionalDetrDecoderLayer, + "attentions": OutputRecorder(ConditionalDetrDecoderSelfAttention, layer_name="self_attn", index=1), + "cross_attentions": OutputRecorder(ConditionalDetrDecoderCrossAttention, layer_name="encoder_attn", index=1), + } + + def __init__(self, config: ConditionalDetrConfig): + super().__init__(config) + self.hidden_size = config.d_model + + self.dropout = config.dropout + self.layerdrop = config.decoder_layerdrop + + self.layers = nn.ModuleList([ConditionalDetrDecoderLayer(config) for _ in range(config.decoder_layers)]) + # in Conditional DETR, the decoder uses layernorm after the last decoder layer output + self.layernorm = nn.LayerNorm(config.d_model) + + # query_scale is the FFN applied on f to generate transformation T + self.query_scale = ConditionalDetrMLPPredictionHead(self.hidden_size, self.hidden_size, self.hidden_size, 2) + self.ref_point_head = ConditionalDetrMLPPredictionHead(self.hidden_size, self.hidden_size, 2, 2) + for layer_id in range(config.decoder_layers - 1): + # Set q_pos_proj to None for layers after the first (only first layer uses query position embeddings) + self.layers[layer_id + 1].encoder_attn.q_pos_proj = None + + # Initialize weights and apply final processing + self.post_init() + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + inputs_embeds=None, + attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + spatial_position_embeddings=None, + object_queries_position_embeddings=None, + **kwargs: Unpack[TransformersKwargs], + ) -> ConditionalDetrDecoderOutput: + r""" + Args: + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): + The query embeddings that are passed into the decoder. + + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on certain queries. Mask values selected in `[0, 1]`: + + - 1 for queries that are **not masked**, + - 0 for queries 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 pixel_values of the encoder. Mask values selected + in `[0, 1]`: + + - 1 for pixels that are real (i.e. **not masked**), + - 0 for pixels that are padding (i.e. **masked**). + + spatial_position_embeddings (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Spatial position embeddings that are added to the queries and keys in each cross-attention layer. + object_queries_position_embeddings (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`): + , *optional*): Position embeddings that are added to the queries and keys in each self-attention layer. + """ + if inputs_embeds is not None: + hidden_states = inputs_embeds + + # expand encoder attention mask + if encoder_hidden_states is not None and encoder_attention_mask is not None: + # [batch_size, seq_len] -> [batch_size, 1, target_seq_len, source_seq_len] + encoder_attention_mask = create_bidirectional_mask( + self.config, + inputs_embeds, + encoder_attention_mask, + ) + + # optional intermediate hidden states + intermediate = () if self.config.auxiliary_loss else None + + reference_points_before_sigmoid = self.ref_point_head( + object_queries_position_embeddings + ) # [num_queries, batch_size, 2] + reference_points = reference_points_before_sigmoid.sigmoid().transpose(0, 1) + obj_center = reference_points[..., :2].transpose(0, 1) + # get sine embedding for the query vector + query_sine_embed_before_transformation = gen_sine_position_embeddings(obj_center, self.config.d_model) + + for idx, decoder_layer in enumerate(self.layers): + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: + continue + if idx == 0: + pos_transformation = 1 + else: + pos_transformation = self.query_scale(hidden_states) + # apply transformation + query_sine_embed = query_sine_embed_before_transformation * pos_transformation + + hidden_states = decoder_layer( + hidden_states, + None, + spatial_position_embeddings, + object_queries_position_embeddings, + query_sine_embed, + encoder_hidden_states, # as a positional argument for gradient checkpointing + encoder_attention_mask=encoder_attention_mask, + is_first=(idx == 0), + **kwargs, + ) + + if self.config.auxiliary_loss: + hidden_states = self.layernorm(hidden_states) + intermediate += (hidden_states,) + + # finally, apply layernorm + hidden_states = self.layernorm(hidden_states) + + # stack intermediate decoder activations + if self.config.auxiliary_loss: + intermediate = torch.stack(intermediate) + + return ConditionalDetrDecoderOutput( + last_hidden_state=hidden_states, + intermediate_hidden_states=intermediate, + reference_points=reference_points, + ) + + +class ConditionalDetrModel(DetrModel): + def __init__(self, config: ConditionalDetrConfig): + super().__init__(config) + self.query_position_embeddings = nn.Embedding(config.num_queries, config.d_model) + + # Initialize weights and apply final processing + self.post_init() + + @auto_docstring + @can_return_tuple + def forward( + self, + pixel_values: torch.FloatTensor, + pixel_mask: torch.LongTensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + encoder_outputs: torch.FloatTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + decoder_inputs_embeds: torch.FloatTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> ConditionalDetrModelOutput: + r""" + decoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_queries)`, *optional*): + Not used by default. Can be used to mask object queries. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing the flattened feature map (output of the backbone + projection layer), you + can choose to directly pass a flattened representation of an image. + decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): + Optionally, instead of initializing the queries with a tensor of zeros, you can choose to directly pass an + embedded representation. + + Examples: + + ```python + >>> from transformers import AutoImageProcessor, AutoModel + >>> from PIL import Image + >>> import requests + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/conditional-detr-resnet-50") + >>> model = AutoModel.from_pretrained("microsoft/conditional-detr-resnet-50") + + >>> # prepare image for the model + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> # forward pass + >>> outputs = model(**inputs) + + >>> # the last hidden states are the final query embeddings of the Transformer decoder + >>> # these are of shape (batch_size, num_queries, hidden_size) + >>> last_hidden_states = outputs.last_hidden_state + >>> list(last_hidden_states.shape) + [1, 300, 256] + ```""" + batch_size, num_channels, height, width = pixel_values.shape + device = pixel_values.device + + if pixel_mask is None: + pixel_mask = torch.ones(((batch_size, height, width)), device=device) + + # First, sent pixel_values + pixel_mask through Backbone to obtain the features + # pixel_values should be of shape (batch_size, num_channels, height, width) + # pixel_mask should be of shape (batch_size, height, width) + features = self.backbone(pixel_values, pixel_mask) + + # get final feature map and downsampled mask + feature_map, mask = features[-1] + + if mask is None: + raise ValueError("Backbone does not return downsampled pixel mask") + + # Second, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default) + projected_feature_map = self.input_projection(feature_map) + + # Generate position embeddings + spatial_position_embeddings = self.position_embedding( + shape=feature_map.shape, device=device, dtype=pixel_values.dtype, mask=mask + ) + + # Third, flatten the feature map of shape NxCxHxW to NxCxHW, and permute it to NxHWxC + # In other words, turn their shape into (batch_size, sequence_length, hidden_size) + flattened_features = projected_feature_map.flatten(2).permute(0, 2, 1) + + flattened_mask = mask.flatten(1) + + # Fourth, sent flattened_features + flattened_mask + spatial_position_embeddings through encoder + # flattened_features is a Tensor of shape (batch_size, height*width, hidden_size) + # flattened_mask is a Tensor of shape (batch_size, height*width) + if encoder_outputs is None: + encoder_outputs = self.encoder( + inputs_embeds=flattened_features, + attention_mask=flattened_mask, + spatial_position_embeddings=spatial_position_embeddings, + **kwargs, + ) + # If the user passed a tuple for encoder_outputs, we wrap it in a BaseModelOutput + 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, + ) + + # Fifth, sent query embeddings through the decoder (which is conditioned on the encoder output) + object_queries_position_embeddings = self.query_position_embeddings.weight.unsqueeze(0).repeat( + batch_size, 1, 1 + ) + queries = torch.zeros_like(object_queries_position_embeddings) + + # decoder outputs consists of (dec_features, dec_hidden, dec_attn) + decoder_outputs = self.decoder( + inputs_embeds=queries, + attention_mask=None, + spatial_position_embeddings=spatial_position_embeddings, + object_queries_position_embeddings=object_queries_position_embeddings, + encoder_hidden_states=encoder_outputs.last_hidden_state, + encoder_attention_mask=flattened_mask, + **kwargs, + ) + + return ConditionalDetrModelOutput( + last_hidden_state=decoder_outputs.last_hidden_state, + 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, + intermediate_hidden_states=decoder_outputs.intermediate_hidden_states, + reference_points=decoder_outputs.reference_points, + ) + + +class ConditionalDetrForObjectDetection(DetrForObjectDetection): + def __init__(self, config: ConditionalDetrConfig): + super().__init__(config) + self.class_labels_classifier = nn.Linear(config.d_model, config.num_labels) + + # taken from https://github.com/Atten4Vis/conditionalDETR/blob/master/models/conditional_detr.py + def _set_aux_loss(self, outputs_class, outputs_coord): + return [{"logits": a, "pred_boxes": b} for a, b in zip(outputs_class[:-1], outputs_coord[:-1])] + + @auto_docstring + @can_return_tuple + def forward( + self, + pixel_values: torch.FloatTensor, + pixel_mask: torch.LongTensor | None = None, + decoder_attention_mask: torch.LongTensor | None = None, + encoder_outputs: torch.FloatTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + decoder_inputs_embeds: torch.FloatTensor | None = None, + labels: list[dict] | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> ConditionalDetrObjectDetectionOutput: + r""" + decoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, num_queries)`, *optional*): + Not used by default. Can be used to mask object queries. + inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): + Optionally, instead of passing the flattened feature map (output of the backbone + projection layer), you + can choose to directly pass a flattened representation of an image. + decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, num_queries, hidden_size)`, *optional*): + Optionally, instead of initializing the queries with a tensor of zeros, you can choose to directly pass an + embedded representation. + labels (`list[Dict]` of len `(batch_size,)`, *optional*): + Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the + following 2 keys: 'class_labels' and 'boxes' (the class labels and bounding boxes of an image in the batch + respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes + in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)`. + + Examples: + + ```python + >>> from transformers import AutoImageProcessor, AutoModelForObjectDetection + >>> from PIL import Image + >>> import requests + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> image_processor = AutoImageProcessor.from_pretrained("microsoft/conditional-detr-resnet-50") + >>> model = AutoModelForObjectDetection.from_pretrained("microsoft/conditional-detr-resnet-50") + + >>> inputs = image_processor(images=image, return_tensors="pt") + + >>> outputs = model(**inputs) + + >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax) + >>> target_sizes = torch.tensor([image.size[::-1]]) + >>> results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[ + ... 0 + ... ] + >>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): + ... box = [round(i, 2) for i in box.tolist()] + ... print( + ... f"Detected {model.config.id2label[label.item()]} with confidence " + ... f"{round(score.item(), 3)} at location {box}" + ... ) + Detected remote with confidence 0.833 at location [38.31, 72.1, 177.63, 118.45] + Detected cat with confidence 0.831 at location [9.2, 51.38, 321.13, 469.0] + Detected cat with confidence 0.804 at location [340.3, 16.85, 642.93, 370.95] + Detected remote with confidence 0.683 at location [334.48, 73.49, 366.37, 190.01] + Detected couch with confidence 0.535 at location [0.52, 1.19, 640.35, 475.1] + ```""" + # First, sent images through CONDITIONAL_DETR base model to obtain encoder + decoder outputs + outputs = self.model( + pixel_values, + pixel_mask=pixel_mask, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + **kwargs, + ) + + sequence_output = outputs[0] + + # class logits + predicted bounding boxes + logits = self.class_labels_classifier(sequence_output) + + reference = outputs.reference_points + reference_before_sigmoid = inverse_sigmoid(reference).transpose(0, 1) + + hs = sequence_output + tmp = self.bbox_predictor(hs) + tmp[..., :2] += reference_before_sigmoid + pred_boxes = tmp.sigmoid() + # pred_boxes = self.bbox_predictor(sequence_output).sigmoid() + + loss, loss_dict, auxiliary_outputs = None, None, None + if labels is not None: + outputs_class, outputs_coord = None, None + if self.config.auxiliary_loss: + outputs_coords = [] + intermediate = outputs.intermediate_hidden_states + outputs_class = self.class_labels_classifier(intermediate) + for lvl in range(intermediate.shape[0]): + tmp = self.bbox_predictor(intermediate[lvl]) + tmp[..., :2] += reference_before_sigmoid + outputs_coord = tmp.sigmoid() + outputs_coords.append(outputs_coord) + outputs_coord = torch.stack(outputs_coords) + loss, loss_dict, auxiliary_outputs = self.loss_function( + logits, labels, self.device, pred_boxes, self.config, outputs_class, outputs_coord + ) + + return ConditionalDetrObjectDetectionOutput( + loss=loss, + loss_dict=loss_dict, + logits=logits, + pred_boxes=pred_boxes, + auxiliary_outputs=auxiliary_outputs, + last_hidden_state=outputs.last_hidden_state, + 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 ConditionalDetrForSegmentation(DetrForSegmentation): + pass + + +__all__ = [ + "ConditionalDetrImageProcessor", + "ConditionalDetrImageProcessorPil", + "ConditionalDetrForObjectDetection", + "ConditionalDetrForSegmentation", + "ConditionalDetrModel", + "ConditionalDetrPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/donut/__init__.py b/third_party/transformers/src/transformers/models/donut/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e1210143546b6deb118bcf489497e6089e7742e5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/donut/__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_donut_swin import * + from .image_processing_donut import * + from .image_processing_pil_donut import * + from .modeling_donut_swin import * + from .processing_donut 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/donut/configuration_donut_swin.py b/third_party/transformers/src/transformers/models/donut/configuration_donut_swin.py new file mode 100644 index 0000000000000000000000000000000000000000..3545ca234c3541c64b3df8a27c626cd418770e4e --- /dev/null +++ b/third_party/transformers/src/transformers/models/donut/configuration_donut_swin.py @@ -0,0 +1,76 @@ +# 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. +"""Donut Swin Transformer model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="naver-clova-ix/donut-base") +@strict +class DonutSwinConfig(PreTrainedConfig): + r""" + window_size (`int`, *optional*, defaults to 7): + Size of windows. + + Example: + + ```python + >>> from transformers import DonutSwinConfig, DonutSwinModel + + >>> # Initializing a Donut naver-clova-ix/donut-base style configuration + >>> configuration = DonutSwinConfig() + + >>> # Randomly initializing a model from the naver-clova-ix/donut-base style configuration + >>> model = DonutSwinModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "donut-swin" + + attribute_map = { + "num_attention_heads": "num_heads", + "num_hidden_layers": "num_layers", + } + + image_size: int | list[int] | tuple[int, int] = 224 + patch_size: int | list[int] | tuple[int, int] = 4 + num_channels: int = 3 + embed_dim: int = 96 + depths: list[int] | tuple[int, ...] = (2, 2, 6, 2) + num_heads: list[int] | tuple[int, ...] = (3, 6, 12, 24) + window_size: int = 7 + mlp_ratio: float = 4.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" + use_absolute_embeddings: bool = False + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-5 + + def __post_init__(self, **kwargs): + self.num_layers = len(self.depths) + # we set the hidden_size attribute in order to make Swin 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)) + super().__post_init__(**kwargs) + + +__all__ = ["DonutSwinConfig"] diff --git a/third_party/transformers/src/transformers/models/donut/convert_donut_to_pytorch.py b/third_party/transformers/src/transformers/models/donut/convert_donut_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..f4172b94d830768d32efd7103e929ffbd9762e49 --- /dev/null +++ b/third_party/transformers/src/transformers/models/donut/convert_donut_to_pytorch.py @@ -0,0 +1,233 @@ +# 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 Donut checkpoints using the original `donut-python` library. URL: https://github.com/clovaai/donut""" + +import argparse + +import torch +from datasets import load_dataset +from donut import DonutModel + +from transformers import ( + DonutImageProcessor, + DonutProcessor, + DonutSwinConfig, + DonutSwinModel, + MBartConfig, + MBartForCausalLM, + VisionEncoderDecoderModel, + XLMRobertaTokenizerFast, +) + + +def get_configs(model): + original_config = model.config + + encoder_config = DonutSwinConfig( + image_size=original_config.input_size, + patch_size=4, + depths=original_config.encoder_layer, + num_heads=[4, 8, 16, 32], + window_size=original_config.window_size, + embed_dim=128, + ) + decoder_config = MBartConfig( + is_decoder=True, + is_encoder_decoder=False, + add_cross_attention=True, + decoder_layers=original_config.decoder_layer, + max_position_embeddings=original_config.max_position_embeddings, + vocab_size=len( + model.decoder.tokenizer + ), # several special tokens are added to the vocab of XLMRobertaTokenizer, see repo on the hub (added_tokens.json) + scale_embedding=True, + add_final_layer_norm=True, + ) + + return encoder_config, decoder_config + + +def rename_key(name): + if "encoder.model" in name: + name = name.replace("encoder.model", "encoder") + if "decoder.model" in name: + name = name.replace("decoder.model", "decoder") + if "patch_embed.proj" in name: + name = name.replace("patch_embed.proj", "embeddings.patch_embeddings.projection") + if "patch_embed.norm" in name: + name = name.replace("patch_embed.norm", "embeddings.norm") + if name.startswith("encoder"): + if "layers" in name: + name = "encoder." + name + if "attn.proj" in name: + name = name.replace("attn.proj", "attention.output.dense") + if "attn" in name and "mask" not in name: + name = name.replace("attn", "attention.self") + if "norm1" in name: + name = name.replace("norm1", "layernorm_before") + if "norm2" in name: + name = name.replace("norm2", "layernorm_after") + if "mlp.fc1" in name: + name = name.replace("mlp.fc1", "intermediate.dense") + if "mlp.fc2" in name: + name = name.replace("mlp.fc2", "output.dense") + + if name == "encoder.norm.weight": + name = "encoder.layernorm.weight" + if name == "encoder.norm.bias": + name = "encoder.layernorm.bias" + + return name + + +def convert_state_dict(orig_state_dict, model): + for key in orig_state_dict.copy(): + val = orig_state_dict.pop(key) + + if "qkv" in key: + key_split = key.split(".") + layer_num = int(key_split[3]) + block_num = int(key_split[5]) + dim = model.encoder.encoder.layers[layer_num].blocks[block_num].attention.self.all_head_size + + if "weight" in key: + orig_state_dict[ + f"encoder.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.query.weight" + ] = val[:dim, :] + orig_state_dict[f"encoder.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.key.weight"] = ( + val[dim : dim * 2, :] + ) + orig_state_dict[ + f"encoder.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.value.weight" + ] = val[-dim:, :] + else: + orig_state_dict[f"encoder.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.query.bias"] = ( + val[:dim] + ) + orig_state_dict[f"encoder.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.key.bias"] = ( + val[dim : dim * 2] + ) + orig_state_dict[f"encoder.encoder.layers.{layer_num}.blocks.{block_num}.attention.self.value.bias"] = ( + val[-dim:] + ) + elif "attn_mask" in key or key in ["encoder.model.norm.weight", "encoder.model.norm.bias"]: + # HuggingFace implementation doesn't use attn_mask buffer + # and model doesn't use final LayerNorms for the encoder + pass + else: + orig_state_dict[rename_key(key)] = val + + return orig_state_dict + + +def convert_donut_checkpoint(model_name, pytorch_dump_folder_path=None, push_to_hub=False): + # load original model + original_model = DonutModel.from_pretrained(model_name).eval() + + # load HuggingFace model + encoder_config, decoder_config = get_configs(original_model) + encoder = DonutSwinModel(encoder_config) + decoder = MBartForCausalLM(decoder_config) + model = VisionEncoderDecoderModel(encoder=encoder, decoder=decoder) + model.eval() + + state_dict = original_model.state_dict() + new_state_dict = convert_state_dict(state_dict, model) + model.load_state_dict(new_state_dict) + + # verify results on scanned document + dataset = load_dataset("hf-internal-testing/example-documents") # no-script + image = dataset["test"][0]["image"].convert("RGB") + + tokenizer = XLMRobertaTokenizerFast.from_pretrained(model_name, from_slow=True) + image_processor = DonutImageProcessor( + do_align_long_axis=original_model.config.align_long_axis, size=original_model.config.input_size[::-1] + ) + processor = DonutProcessor(image_processor, tokenizer) + pixel_values = processor(image, return_tensors="pt").pixel_values + + if model_name == "naver-clova-ix/donut-base-finetuned-docvqa": + task_prompt = "{user_input}" + question = "When is the coffee break?" + task_prompt = task_prompt.replace("{user_input}", question) + elif model_name == "naver-clova-ix/donut-base-finetuned-rvlcdip": + task_prompt = "" + elif model_name in [ + "naver-clova-ix/donut-base-finetuned-cord-v1", + "naver-clova-ix/donut-base-finetuned-cord-v1-2560", + ]: + task_prompt = "" + elif model_name == "naver-clova-ix/donut-base-finetuned-cord-v2": + task_prompt = "s_cord-v2>" + elif model_name == "naver-clova-ix/donut-base-finetuned-zhtrainticket": + task_prompt = "" + elif model_name in ["naver-clova-ix/donut-proto", "naver-clova-ix/donut-base"]: + # use a random prompt + task_prompt = "hello world" + else: + raise ValueError("Model name not supported") + prompt_tensors = original_model.decoder.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt")[ + "input_ids" + ] + + original_patch_embed = original_model.encoder.model.patch_embed(pixel_values) + patch_embeddings, _ = model.encoder.embeddings(pixel_values) + assert torch.allclose(original_patch_embed, patch_embeddings, atol=1e-3) + + # verify encoder hidden states + original_last_hidden_state = original_model.encoder(pixel_values) + last_hidden_state = model.encoder(pixel_values).last_hidden_state + assert torch.allclose(original_last_hidden_state, last_hidden_state, atol=1e-2) + + # verify decoder hidden states + original_logits = original_model(pixel_values, prompt_tensors, None).logits + logits = model(pixel_values, decoder_input_ids=prompt_tensors).logits + assert torch.allclose(original_logits, logits, 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: + model.push_to_hub("nielsr/" + model_name.split("/")[-1], commit_message="Update model") + processor.push_to_hub("nielsr/" + model_name.split("/")[-1], commit_message="Update model") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--model_name", + default="naver-clova-ix/donut-base-finetuned-docvqa", + required=False, + type=str, + help="Name of the original model you'd like to convert.", + ) + parser.add_argument( + "--pytorch_dump_folder_path", + default=None, + required=False, + 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 and processor to the Hugging Face hub.", + ) + + args = parser.parse_args() + convert_donut_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub) diff --git a/third_party/transformers/src/transformers/models/donut/image_processing_donut.py b/third_party/transformers/src/transformers/models/donut/image_processing_donut.py new file mode 100644 index 0000000000000000000000000000000000000000..9e4cb28e2be94ceaf88fc5185fe5e406ab6653a5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/donut/image_processing_donut.py @@ -0,0 +1,206 @@ +# 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 Donut.""" + +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 ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ImageInput, + PILImageResampling, + SizeDict, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring + + +class DonutImageProcessorKwargs(ImagesKwargs, total=False): + r""" + do_thumbnail (`bool`, *optional*, defaults to `self.do_thumbnail`): + Whether to resize the image using thumbnail method. + do_align_long_axis (`bool`, *optional*, defaults to `self.do_align_long_axis`): + Whether to align the long axis of the image with the long axis of `size` by rotating by 90 degrees. + """ + + do_thumbnail: bool + do_align_long_axis: bool + + +@auto_docstring +class DonutImageProcessor(TorchvisionBackend): + """Torchvision backend for Donut with align_long_axis, thumbnail, and pad_image.""" + + valid_kwargs = DonutImageProcessorKwargs + + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 2560, "width": 1920} + do_resize = True + do_rescale = True + do_normalize = True + do_thumbnail = True + do_align_long_axis = False + do_pad = True + + def __init__(self, **kwargs: Unpack[DonutImageProcessorKwargs]): + size = kwargs.pop("size", None) + if isinstance(size, (tuple, list)): + size = size[::-1] + if size is not None: + kwargs["size"] = size + super().__init__(**kwargs) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + **kwargs: Unpack[DonutImageProcessorKwargs], + ) -> BatchFeature: + kwargs = dict(kwargs) + if "size" in kwargs: + size = kwargs["size"] + if isinstance(size, (tuple, list)): + kwargs["size"] = size[::-1] + return super().preprocess(images, **kwargs) + + def align_long_axis( + self, + image: "torch.Tensor", + size: SizeDict, + ) -> "torch.Tensor": + """Align the long axis of the image to the longest axis of the specified size.""" + input_height, input_width = image.shape[-2:] + output_height, output_width = size.height, size.width + + if (output_width < output_height and input_width > input_height) or ( + output_width > output_height and input_width < input_height + ): + height_dim, width_dim = image.dim() - 2, image.dim() - 1 + image = torch.rot90(image, 3, dims=[height_dim, width_dim]) + + return image + + def pad_image( + self, + image: "torch.Tensor", + size: SizeDict, + random_padding: bool = False, + ) -> "torch.Tensor": + """Pad the image to the specified size.""" + output_height, output_width = size.height, size.width + input_height, input_width = image.shape[-2:] + + delta_width = output_width - input_width + delta_height = output_height - input_height + + if random_padding: + pad_top = torch.randint(0, delta_height + 1, ()).item() + pad_left = torch.randint(0, delta_width + 1, ()).item() + else: + pad_top = delta_height // 2 + pad_left = delta_width // 2 + + pad_bottom = delta_height - pad_top + pad_right = delta_width - pad_left + + padding = (pad_left, pad_top, pad_right, pad_bottom) + return tvF.pad(image, padding) + + def thumbnail( + self, + image: "torch.Tensor", + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None, + **kwargs, + ) -> "torch.Tensor": + """Resize the image to make a thumbnail.""" + input_height, input_width = image.shape[-2:] + output_height, output_width = size.height, size.width + + height = min(input_height, output_height) + width = min(input_width, output_width) + + if height == input_height and width == input_width: + return image + + if input_height > input_width: + width = int(input_width * height / input_height) + elif input_width > input_height: + height = int(input_height * width / input_width) + + return super().resize( + image, + size=SizeDict(width=width, height=height), + 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, + do_thumbnail: bool = True, + do_align_long_axis: bool = False, + **kwargs, + ) -> BatchFeature: + """Custom preprocessing for Donut.""" + 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_align_long_axis: + stacked_images = self.align_long_axis(stacked_images, size) + if do_resize: + shortest_edge = min(size.height, size.width) + stacked_images = self.resize(stacked_images, SizeDict(shortest_edge=shortest_edge), resample) + if do_thumbnail: + stacked_images = self.thumbnail(stacked_images, size, resample) + if do_pad: + stacked_images = self.pad_image(stacked_images, size, random_padding=False) + 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__ = ["DonutImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/donut/image_processing_pil_donut.py b/third_party/transformers/src/transformers/models/donut/image_processing_pil_donut.py new file mode 100644 index 0000000000000000000000000000000000000000..b2865b556bc8f3d73695bdb607e6b8c8ca56a88f --- /dev/null +++ b/third_party/transformers/src/transformers/models/donut/image_processing_pil_donut.py @@ -0,0 +1,212 @@ +# 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 Donut.""" + +import numpy as np + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ImageInput, + PILImageResampling, + SizeDict, + get_image_size, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring + + +# Adapted from transformers.models.donut.image_processing_donut.DonutImageProcessorKwargs +class DonutImageProcessorKwargs(ImagesKwargs, total=False): + r""" + do_thumbnail (`bool`, *optional*, defaults to `self.do_thumbnail`): + Whether to resize the image using thumbnail method. + do_align_long_axis (`bool`, *optional*, defaults to `self.do_align_long_axis`): + Whether to align the long axis of the image with the long axis of `size` by rotating by 90 degrees. + """ + + do_thumbnail: bool + do_align_long_axis: bool + + +@auto_docstring +class DonutImageProcessorPil(PilBackend): + """PIL backend for Donut with align_long_axis, thumbnail, and pad_image.""" + + valid_kwargs = DonutImageProcessorKwargs + + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 2560, "width": 1920} + do_resize = True + do_rescale = True + do_normalize = True + do_thumbnail = True + do_align_long_axis = False + do_pad = True + + def __init__(self, **kwargs: Unpack[DonutImageProcessorKwargs]): + size = kwargs.pop("size", None) + if isinstance(size, (tuple, list)): + size = size[::-1] + if size is not None: + kwargs["size"] = size + super().__init__(**kwargs) + + @auto_docstring + def preprocess( + self, + images: ImageInput, + **kwargs: Unpack[DonutImageProcessorKwargs], + ) -> BatchFeature: + kwargs = dict(kwargs) + if "size" in kwargs: + size = kwargs["size"] + if isinstance(size, (tuple, list)): + kwargs["size"] = size[::-1] + return super().preprocess(images, **kwargs) + + def align_long_axis( + self, + image: np.ndarray, + size: SizeDict, + ) -> np.ndarray: + """Align the long axis of the image to the longest axis of the specified size.""" + from ...image_utils import ChannelDimension + + input_height, input_width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + output_height, output_width = size.height, size.width + + if (output_width < output_height and input_width > input_height) or ( + output_width > output_height and input_width < input_height + ): + image = np.rot90(image, 3, axes=(1, 2)) + + return image + + def pad_image( + self, + image: np.ndarray, + size: SizeDict, + random_padding: bool = False, + ) -> np.ndarray: + """Pad the image to the specified size.""" + from ...image_transforms import PaddingMode + from ...image_transforms import pad as np_pad + from ...image_utils import ChannelDimension + + output_height, output_width = size.height, size.width + input_height, input_width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + + delta_width = output_width - input_width + delta_height = output_height - input_height + + if random_padding: + pad_top = int(np.random.randint(low=0, high=delta_height + 1)) + pad_left = int(np.random.randint(low=0, high=delta_width + 1)) + else: + pad_top = delta_height // 2 + pad_left = delta_width // 2 + + pad_bottom = delta_height - pad_top + pad_right = delta_width - pad_left + + # pad() expects (height_pad, width_pad) and adds channel dimension + padding = ((pad_top, pad_bottom), (pad_left, pad_right)) + return np_pad( + image, + padding, + mode=PaddingMode.CONSTANT, + constant_values=0, + data_format=ChannelDimension.FIRST, + input_data_format=ChannelDimension.FIRST, + ) + + def thumbnail( + self, + image: np.ndarray, + size: SizeDict, + resample: "PILImageResampling | None" = None, + **kwargs, + ) -> np.ndarray: + """Resize the image to make a thumbnail.""" + from ...image_utils import ChannelDimension + + input_height, input_width = get_image_size(image, channel_dim=ChannelDimension.FIRST) + output_height, output_width = size.height, size.width + + height = min(input_height, output_height) + width = min(input_width, output_width) + + if height == input_height and width == input_width: + return image + + if input_height > input_width: + width = int(input_width * height / input_height) + elif input_width > input_height: + height = int(input_height * width / input_width) + + return self.resize( + image, + size=SizeDict(width=width, height=height), + resample=resample or PILImageResampling.BICUBIC, + ) + + 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, + do_thumbnail: bool = True, + do_align_long_axis: bool = False, + **kwargs, + ) -> BatchFeature: + """Custom preprocessing for Donut.""" + processed_images = [] + for image in images: + if do_align_long_axis: + image = self.align_long_axis(image, size) + if do_resize: + shortest_edge = min(size.height, size.width) + image = self.resize(image, SizeDict(shortest_edge=shortest_edge), resample) + if do_thumbnail: + image = self.thumbnail(image, size, resample) + if do_pad: + image = self.pad_image(image, size, random_padding=False) + 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__ = ["DonutImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/donut/modeling_donut_swin.py b/third_party/transformers/src/transformers/models/donut/modeling_donut_swin.py new file mode 100644 index 0000000000000000000000000000000000000000..b3521146938a16c4d4bc27acfe030dd398482cae --- /dev/null +++ b/third_party/transformers/src/transformers/models/donut/modeling_donut_swin.py @@ -0,0 +1,967 @@ +# 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. +"""PyTorch Donut Swin Transformer model. + +This implementation is identical to a regular Swin Transformer, without final layer norm on top of the final hidden +states.""" + +import collections.abc +import math +from dataclasses import dataclass + +import torch +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_utils import PreTrainedModel +from ...utils import ModelOutput, auto_docstring, logging, torch_int +from .configuration_donut_swin import DonutSwinConfig + + +logger = logging.get_logger(__name__) + + +@dataclass +@auto_docstring( + custom_intro=""" + DonutSwin encoder's outputs, with potential hidden states and attentions. + """ +) +# Copied from transformers.models.swin.modeling_swin.SwinEncoderOutput with Swin->DonutSwin +class DonutSwinEncoderOutput(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=""" + DonutSwin model's outputs that also contains a pooling of the last hidden states. + """ +) +# Copied from transformers.models.swin.modeling_swin.SwinModelOutput with Swin->DonutSwin +class DonutSwinModelOutput(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=""" + DonutSwin outputs for image classification. + """ +) +# Copied from transformers.models.swin.modeling_swin.SwinImageClassifierOutput with Swin->DonutSwin +class DonutSwinImageClassifierOutput(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 + + +# Copied from transformers.models.swin.modeling_swin.window_partition +def window_partition(input_feature, window_size): + """ + Partitions the given input into windows. + """ + batch_size, height, width, num_channels = input_feature.shape + input_feature = input_feature.view( + batch_size, height // window_size, window_size, width // window_size, window_size, num_channels + ) + windows = input_feature.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size, window_size, num_channels) + return windows + + +# Copied from transformers.models.swin.modeling_swin.window_reverse +def window_reverse(windows, window_size, height, width): + """ + Merges windows to produce higher resolution features. + """ + 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 + + +# Copied from transformers.models.swin.modeling_swin.SwinEmbeddings with Swin->DonutSwin +class DonutSwinEmbeddings(nn.Module): + """ + Construct the patch and position embeddings. Optionally, also the mask token. + """ + + def __init__(self, config, use_mask_token=False): + super().__init__() + + self.patch_embeddings = DonutSwinPatchEmbeddings(config) + num_patches = self.patch_embeddings.num_patches + self.patch_grid = self.patch_embeddings.grid_size + self.mask_token = nn.Parameter(torch.zeros(1, 1, config.embed_dim)) if use_mask_token else None + + if config.use_absolute_embeddings: + self.position_embeddings = nn.Parameter(torch.zeros(1, num_patches + 1, config.embed_dim)) + else: + self.position_embeddings = None + + self.norm = nn.LayerNorm(config.embed_dim) + 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.FloatTensor | None, + bool_masked_pos: torch.BoolTensor | None = None, + interpolate_pos_encoding: bool = False, + ) -> tuple[torch.Tensor]: + _, num_channels, height, width = pixel_values.shape + embeddings, output_dimensions = self.patch_embeddings(pixel_values) + embeddings = self.norm(embeddings) + batch_size, seq_len, _ = embeddings.size() + + if bool_masked_pos is not None: + mask_tokens = self.mask_token.expand(batch_size, seq_len, -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 + + if self.position_embeddings is not None: + 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, output_dimensions + + +# Copied from transformers.models.swin.modeling_swin.SwinPatchEmbeddings with Swin->DonutSwin +class DonutSwinPatchEmbeddings(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.embed_dim + 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.grid_size = (image_size[0] // patch_size[0], image_size[1] // patch_size[1]) + + self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size) + + def maybe_pad(self, pixel_values, height, width): + if width % self.patch_size[1] != 0: + pad_values = (0, self.patch_size[1] - width % self.patch_size[1]) + pixel_values = nn.functional.pad(pixel_values, pad_values) + if height % self.patch_size[0] != 0: + pad_values = (0, 0, 0, self.patch_size[0] - height % self.patch_size[0]) + pixel_values = nn.functional.pad(pixel_values, pad_values) + return pixel_values + + def forward(self, pixel_values: torch.FloatTensor | None) -> tuple[torch.Tensor, tuple[int]]: + _, num_channels, height, width = pixel_values.shape + # pad the input to be divisible by self.patch_size, if needed + pixel_values = self.maybe_pad(pixel_values, height, width) + embeddings = self.projection(pixel_values) + _, _, height, width = embeddings.shape + output_dimensions = (height, width) + embeddings = embeddings.flatten(2).transpose(1, 2) + + return embeddings, output_dimensions + + +# Copied from transformers.models.swin.modeling_swin.SwinPatchMerging +class DonutSwinPatchMerging(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 + + +# 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.swin.modeling_swin.SwinDropPath +class DonutSwinDropPath(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}" + + +# Copied from transformers.models.swin.modeling_swin.SwinSelfAttention with Swin->DonutSwin +class DonutSwinSelfAttention(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 DonutSwinModel 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 +class DonutSwinSelfOutput(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->DonutSwin +class DonutSwinAttention(nn.Module): + def __init__(self, config, dim, num_heads, window_size): + super().__init__() + self.self = DonutSwinSelfAttention(config, dim, num_heads, window_size) + self.output = DonutSwinSelfOutput(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 +class DonutSwinIntermediate(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 +class DonutSwinOutput(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 Swin->DonutSwin +class DonutSwinLayer(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 = DonutSwinAttention(config, dim, num_heads, window_size=self.window_size) + self.drop_path = DonutSwinDropPath(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 = DonutSwinIntermediate(config, dim) + self.output = DonutSwinOutput(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->DonutSwin +class DonutSwinStage(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( + [ + DonutSwinLayer( + 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.SwinEncoder with Swin->DonutSwin +class DonutSwinEncoder(nn.Module): + def __init__(self, config, grid_size): + super().__init__() + self.num_layers = 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.layers = nn.ModuleList( + [ + DonutSwinStage( + config=config, + dim=int(config.embed_dim * 2**i_layer), + input_resolution=(grid_size[0] // (2**i_layer), grid_size[1] // (2**i_layer)), + depth=config.depths[i_layer], + num_heads=config.num_heads[i_layer], + drop_path=dpr[sum(config.depths[:i_layer]) : sum(config.depths[: i_layer + 1])], + downsample=DonutSwinPatchMerging if (i_layer < self.num_layers - 1) else None, + ) + for i_layer in range(self.num_layers) + ] + ) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + input_dimensions: tuple[int, int], + 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 | DonutSwinEncoderOutput: + 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: + batch_size, _, hidden_size = hidden_states.shape + # rearrange b (h w) c -> b c h w + 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): + 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 b (h w) c -> b c h w + # 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 b (h w) c -> b c h w + 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:] + + if not return_dict: + return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None) + + return DonutSwinEncoderOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + reshaped_hidden_states=all_reshaped_hidden_states, + ) + + +@auto_docstring +# Copied from transformers.models.swin.modeling_swin.SwinPreTrainedModel with Swin->DonutSwin,swin->donut +class DonutSwinPreTrainedModel(PreTrainedModel): + config: DonutSwinConfig + base_model_prefix = "donut" + main_input_name = "pixel_values" + input_modalities = ("image",) + supports_gradient_checkpointing = True + _no_split_modules = ["DonutSwinStage"] + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + if isinstance(module, DonutSwinEmbeddings): + if module.mask_token is not None: + init.zeros_(module.mask_token) + if module.position_embeddings is not None: + init.zeros_(module.position_embeddings) + elif isinstance(module, DonutSwinSelfAttention): + init.zeros_(module.relative_position_bias_table) + init.copy_(module.relative_position_index, module.create_relative_position_index()) + + +@auto_docstring +class DonutSwinModel(DonutSwinPreTrainedModel): + def __init__(self, config, add_pooling_layer=True, use_mask_token=False): + r""" + add_pooling_layer (bool, *optional*, defaults to `True`): + Whether to add a pooling layer + use_mask_token (`bool`, *optional*, defaults to `False`): + Whether to use a mask token for masked image modeling. + """ + super().__init__(config) + self.config = config + self.num_layers = len(config.depths) + self.num_features = int(config.embed_dim * 2 ** (self.num_layers - 1)) + + self.embeddings = DonutSwinEmbeddings(config, use_mask_token=use_mask_token) + self.encoder = DonutSwinEncoder(config, self.embeddings.patch_grid) + + 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, + bool_masked_pos: torch.BoolTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + interpolate_pos_encoding: bool = False, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | DonutSwinModelOutput: + r""" + bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, 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, input_dimensions = self.embeddings( + pixel_values, bool_masked_pos=bool_masked_pos, interpolate_pos_encoding=interpolate_pos_encoding + ) + + encoder_outputs = self.encoder( + embedding_output, + input_dimensions, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = encoder_outputs[0] + + pooled_output = None + if self.pooler is not None: + pooled_output = self.pooler(sequence_output.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 DonutSwinModelOutput( + 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=""" + DonutSwin 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. + + + + Note that it's possible to fine-tune DonutSwin on higher resolution images than the ones it has been trained on, by + setting `interpolate_pos_encoding` to `True` in the forward of the model. This will interpolate the pre-trained + position embeddings to the higher resolution. + + + """ +) +# Copied from transformers.models.swin.modeling_swin.SwinForImageClassification with Swin->DonutSwin,swin->donut +class DonutSwinForImageClassification(DonutSwinPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.num_labels = config.num_labels + self.donut = DonutSwinModel(config) + + # Classifier head + self.classifier = ( + nn.Linear(self.donut.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, + interpolate_pos_encoding: bool = False, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | DonutSwinImageClassifierOutput: + 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.donut( + pixel_values, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + interpolate_pos_encoding=interpolate_pos_encoding, + 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 DonutSwinImageClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + reshaped_hidden_states=outputs.reshaped_hidden_states, + ) + + +__all__ = ["DonutSwinModel", "DonutSwinPreTrainedModel", "DonutSwinForImageClassification"] diff --git a/third_party/transformers/src/transformers/models/donut/processing_donut.py b/third_party/transformers/src/transformers/models/donut/processing_donut.py new file mode 100644 index 0000000000000000000000000000000000000000..1db500723f1ccff24215e95c42853897d78597f7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/donut/processing_donut.py @@ -0,0 +1,135 @@ +# 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 Donut. +""" + +import re + +from ...image_utils import ImageInput +from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import auto_docstring, logging + + +class DonutProcessorKwargs(ProcessingKwargs, total=False): + _defaults = {} + + +logger = logging.get_logger(__name__) + + +@auto_docstring +class DonutProcessor(ProcessorMixin): + def __init__(self, image_processor=None, tokenizer=None, **kwargs): + super().__init__(image_processor, tokenizer) + + @auto_docstring + def __call__( + self, + images: ImageInput | None = None, + text: str | list[str] | TextInput | PreTokenizedInput | None = None, + **kwargs: Unpack[DonutProcessorKwargs], + ): + if images is None and text is None: + raise ValueError("You need to specify either an `images` or `text` input to process.") + + output_kwargs = self._merge_kwargs( + DonutProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + if images is not None: + inputs = self.image_processor(images, **output_kwargs["images_kwargs"]) + if text is not None: + if images is not None: + output_kwargs["text_kwargs"].setdefault("add_special_tokens", False) + encodings = self.tokenizer(text, **output_kwargs["text_kwargs"]) + + if text is None: + return inputs + elif images is None: + return encodings + else: + inputs["labels"] = encodings["input_ids"] # for BC + inputs["input_ids"] = encodings["input_ids"] + return inputs + + @property + def model_input_names(self): + image_processor_input_names = self.image_processor.model_input_names + + return list(image_processor_input_names + ["input_ids", "labels"]) + + def token2json(self, tokens, is_inner_value=False, added_vocab=None): + """ + Convert a (generated) token sequence into an ordered JSON format. + """ + if added_vocab is None: + added_vocab = self.tokenizer.get_added_vocab() + + output = {} + + while tokens: + # We want r"" but without ReDOS risk, so do it manually in two parts + potential_start = re.search(r"" not in start_token: + break + start_token = start_token[: start_token.index(">") + 1] + key = start_token[len("")] + key_escaped = re.escape(key) + + end_token = re.search(rf"", tokens, re.IGNORECASE) + if end_token is None: + tokens = tokens.replace(start_token, "") + else: + end_token = end_token.group() + start_token_escaped = re.escape(start_token) + end_token_escaped = re.escape(end_token) + content = re.search( + f"{start_token_escaped}(.*?){end_token_escaped}", tokens, re.IGNORECASE | re.DOTALL + ) + if content is not None: + content = content.group(1).strip() + if r""): + leaf = leaf.strip() + if leaf in added_vocab and leaf[0] == "<" and leaf[-2:] == "/>": + leaf = leaf[1:-2] # for categorical special tokens + output[key].append(leaf) + if len(output[key]) == 1: + output[key] = output[key][0] + + tokens = tokens[tokens.find(end_token) + len(end_token) :].strip() + if tokens[:6] == r"": # non-leaf nodes + return [output] + self.token2json(tokens[6:], is_inner_value=True, added_vocab=added_vocab) + + if output: + return [output] if is_inner_value else output + else: + return [] if is_inner_value else {"text_sequence": tokens} + + +__all__ = ["DonutProcessor"] diff --git a/third_party/transformers/src/transformers/models/ibert/__init__.py b/third_party/transformers/src/transformers/models/ibert/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cf34ec43ac1014d8c153b3aa259e394fc7b73570 --- /dev/null +++ b/third_party/transformers/src/transformers/models/ibert/__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_ibert import * + from .modeling_ibert 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/ibert/configuration_ibert.py b/third_party/transformers/src/transformers/models/ibert/configuration_ibert.py new file mode 100644 index 0000000000000000000000000000000000000000..07e9a87b6a0084233a60a2aca421ed3f743e7dfd --- /dev/null +++ b/third_party/transformers/src/transformers/models/ibert/configuration_ibert.py @@ -0,0 +1,61 @@ +# Copyright 2021 The I-BERT Authors (Sehoon Kim, Amir Gholami, Zhewei Yao, +# Michael Mahoney, Kurt Keutzer - UC Berkeley) and The HuggingFace Inc. team. +# Copyright (c) 20121, 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. +"""I-BERT configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="kssteven/ibert-roberta-base") +@strict +class IBertConfig(PreTrainedConfig): + r""" + type_vocab_size (`int`, *optional*, defaults to 2): + The vocabulary size of the `token_type_ids` passed when calling [`IBertModel`] + quant_mode (`bool`, *optional*, defaults to `False`): + Whether to quantize the model or not. + force_dequant (`str`, *optional*, defaults to `"none"`): + Force dequantize specific nonlinear layer. Dequantized layers are then executed with full precision. + `"none"`, `"gelu"`, `"softmax"`, `"layernorm"` and `"nonlinear"` are supported. As default, it is set as + `"none"`, which does not dequantize any layers. Please specify `"gelu"`, `"softmax"`, or `"layernorm"` to + dequantize GELU, Softmax, or LayerNorm, respectively. `"nonlinear"` will dequantize all nonlinear layers, + i.e., GELU, Softmax, and LayerNorm. + """ + + model_type = "ibert" + + 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 = 1 + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 2 + quant_mode: bool = False + force_dequant: str = "none" + + +__all__ = ["IBertConfig"] diff --git a/third_party/transformers/src/transformers/models/ibert/modeling_ibert.py b/third_party/transformers/src/transformers/models/ibert/modeling_ibert.py new file mode 100644 index 0000000000000000000000000000000000000000..58c0dd37d3c18cf20f2522028ab5f39eb928bf24 --- /dev/null +++ b/third_party/transformers/src/transformers/models/ibert/modeling_ibert.py @@ -0,0 +1,1201 @@ +# Copyright 2021 The I-BERT Authors (Sehoon Kim, Amir Gholami, Zhewei Yao, +# Michael Mahoney, Kurt Keutzer - UC Berkeley) and The HuggingFace Inc. team. +# Copyright (c) 20121, 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 I-BERT model.""" + +import math + +import torch +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from ... import initialization as init +from ...activations import gelu +from ...modeling_outputs import ( + BaseModelOutputWithPastAndCrossAttentions, + BaseModelOutputWithPoolingAndCrossAttentions, + MaskedLMOutput, + MultipleChoiceModelOutput, + QuestionAnsweringModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from ...modeling_utils import PreTrainedModel +from ...utils import auto_docstring, logging +from .configuration_ibert import IBertConfig +from .quant_modules import IntGELU, IntLayerNorm, IntSoftmax, QuantAct, QuantEmbedding, QuantLinear + + +logger = logging.get_logger(__name__) + + +class IBertEmbeddings(nn.Module): + """ + Same as BertEmbeddings with a tiny tweak for positional embeddings indexing. + """ + + def __init__(self, config): + super().__init__() + self.quant_mode = config.quant_mode + self.embedding_bit = 8 + self.embedding_act_bit = 16 + self.act_bit = 8 + self.ln_input_bit = 22 + self.ln_output_bit = 32 + + self.word_embeddings = QuantEmbedding( + config.vocab_size, + config.hidden_size, + padding_idx=config.pad_token_id, + weight_bit=self.embedding_bit, + quant_mode=self.quant_mode, + ) + self.token_type_embeddings = QuantEmbedding( + config.type_vocab_size, config.hidden_size, weight_bit=self.embedding_bit, quant_mode=self.quant_mode + ) + + # 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 + ) + + # End copy + self.padding_idx = config.pad_token_id + self.position_embeddings = QuantEmbedding( + config.max_position_embeddings, + config.hidden_size, + padding_idx=self.padding_idx, + weight_bit=self.embedding_bit, + quant_mode=self.quant_mode, + ) + + # Integer-only addition between embeddings + self.embeddings_act1 = QuantAct(self.embedding_act_bit, quant_mode=self.quant_mode) + self.embeddings_act2 = QuantAct(self.embedding_act_bit, quant_mode=self.quant_mode) + + self.LayerNorm = IntLayerNorm( + config.hidden_size, + eps=config.layer_norm_eps, + output_bit=self.ln_output_bit, + quant_mode=self.quant_mode, + force_dequant=config.force_dequant, + ) + self.output_activation = QuantAct(self.act_bit, quant_mode=self.quant_mode) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward( + self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0 + ): + 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 = create_position_ids_from_input_ids( + input_ids, self.padding_idx, past_key_values_length + ).to(input_ids.device) + else: + position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds) + + if input_ids is not None: + input_shape = input_ids.size() + else: + input_shape = inputs_embeds.size()[:-1] + + if token_type_ids is None: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device) + + if inputs_embeds is None: + inputs_embeds, inputs_embeds_scaling_factor = self.word_embeddings(input_ids) + else: + inputs_embeds_scaling_factor = None + token_type_embeddings, token_type_embeddings_scaling_factor = self.token_type_embeddings(token_type_ids) + + embeddings, embeddings_scaling_factor = self.embeddings_act1( + inputs_embeds, + inputs_embeds_scaling_factor, + identity=token_type_embeddings, + identity_scaling_factor=token_type_embeddings_scaling_factor, + ) + + position_embeddings, position_embeddings_scaling_factor = self.position_embeddings(position_ids) + embeddings, embeddings_scaling_factor = self.embeddings_act1( + embeddings, + embeddings_scaling_factor, + identity=position_embeddings, + identity_scaling_factor=position_embeddings_scaling_factor, + ) + + embeddings, embeddings_scaling_factor = self.LayerNorm(embeddings, embeddings_scaling_factor) + embeddings = self.dropout(embeddings) + embeddings, embeddings_scaling_factor = self.output_activation(embeddings, embeddings_scaling_factor) + return embeddings, embeddings_scaling_factor + + def create_position_ids_from_inputs_embeds(self, inputs_embeds): + """ + 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( + self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device + ) + return position_ids.unsqueeze(0).expand(input_shape) + + +class IBertSelfAttention(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.quant_mode = config.quant_mode + self.weight_bit = 8 + self.bias_bit = 32 + self.act_bit = 8 + + 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 + + # Q, K, V Linear layers + self.query = QuantLinear( + config.hidden_size, + self.all_head_size, + bias=True, + weight_bit=self.weight_bit, + bias_bit=self.bias_bit, + quant_mode=self.quant_mode, + per_channel=True, + ) + self.key = QuantLinear( + config.hidden_size, + self.all_head_size, + bias=True, + weight_bit=self.weight_bit, + bias_bit=self.bias_bit, + quant_mode=self.quant_mode, + per_channel=True, + ) + self.value = QuantLinear( + config.hidden_size, + self.all_head_size, + bias=True, + weight_bit=self.weight_bit, + bias_bit=self.bias_bit, + quant_mode=self.quant_mode, + per_channel=True, + ) + + # Requantization (32bit -> 8bit) for Q, K, V activations + self.query_activation = QuantAct(self.act_bit, quant_mode=self.quant_mode) + self.key_activation = QuantAct(self.act_bit, quant_mode=self.quant_mode) + self.value_activation = QuantAct(self.act_bit, quant_mode=self.quant_mode) + self.output_activation = QuantAct(self.act_bit, quant_mode=self.quant_mode) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + self.softmax = IntSoftmax(self.act_bit, quant_mode=self.quant_mode, force_dequant=config.force_dequant) + + def forward( + self, + hidden_states, + hidden_states_scaling_factor, + attention_mask=None, + output_attentions=False, + ): + # Projection + mixed_query_layer, mixed_query_layer_scaling_factor = self.query(hidden_states, hidden_states_scaling_factor) + mixed_key_layer, mixed_key_layer_scaling_factor = self.key(hidden_states, hidden_states_scaling_factor) + mixed_value_layer, mixed_value_layer_scaling_factor = self.value(hidden_states, hidden_states_scaling_factor) + + # Requantization + query_layer, query_layer_scaling_factor = self.query_activation( + mixed_query_layer, mixed_query_layer_scaling_factor + ) + key_layer, key_layer_scaling_factor = self.key_activation(mixed_key_layer, mixed_key_layer_scaling_factor) + value_layer, value_layer_scaling_factor = self.value_activation( + mixed_value_layer, mixed_value_layer_scaling_factor + ) + + # Transpose + batch_size, seq_length, _ = hidden_states.shape + query_layer = query_layer.view(batch_size, -1, self.num_attention_heads, self.attention_head_size).transpose( + 1, 2 + ) + key_layer = key_layer.view(batch_size, -1, self.num_attention_heads, self.attention_head_size).transpose(1, 2) + value_layer = value_layer.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)) + scale = math.sqrt(self.attention_head_size) + attention_scores = attention_scores / scale + if self.quant_mode: + attention_scores_scaling_factor = query_layer_scaling_factor * key_layer_scaling_factor / scale + else: + attention_scores_scaling_factor = None + + if attention_mask is not None: + # Apply the attention mask is (precomputed for all layers in IBertModel forward() function) + attention_scores = attention_scores + attention_mask + + # Normalize the attention scores to probabilities. + attention_probs, attention_probs_scaling_factor = self.softmax( + attention_scores, attention_scores_scaling_factor + ) + + # 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) + if attention_probs_scaling_factor is not None: + context_layer_scaling_factor = attention_probs_scaling_factor * value_layer_scaling_factor + else: + context_layer_scaling_factor = None + + 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) + + # requantization: 32-bit -> 8-bit + context_layer, context_layer_scaling_factor = self.output_activation( + context_layer, context_layer_scaling_factor + ) + + outputs = (context_layer, attention_probs) if output_attentions else (context_layer,) + output_scaling_factor = ( + (context_layer_scaling_factor, attention_probs_scaling_factor) + if output_attentions + else (context_layer_scaling_factor,) + ) + + return outputs, output_scaling_factor + + +class IBertSelfOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.quant_mode = config.quant_mode + self.act_bit = 8 + self.weight_bit = 8 + self.bias_bit = 32 + self.ln_input_bit = 22 + self.ln_output_bit = 32 + + self.dense = QuantLinear( + config.hidden_size, + config.hidden_size, + bias=True, + weight_bit=self.weight_bit, + bias_bit=self.bias_bit, + quant_mode=self.quant_mode, + per_channel=True, + ) + self.ln_input_act = QuantAct(self.ln_input_bit, quant_mode=self.quant_mode) + self.LayerNorm = IntLayerNorm( + config.hidden_size, + eps=config.layer_norm_eps, + output_bit=self.ln_output_bit, + quant_mode=self.quant_mode, + force_dequant=config.force_dequant, + ) + self.output_activation = QuantAct(self.act_bit, quant_mode=self.quant_mode) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states, hidden_states_scaling_factor, input_tensor, input_tensor_scaling_factor): + hidden_states, hidden_states_scaling_factor = self.dense(hidden_states, hidden_states_scaling_factor) + hidden_states = self.dropout(hidden_states) + hidden_states, hidden_states_scaling_factor = self.ln_input_act( + hidden_states, + hidden_states_scaling_factor, + identity=input_tensor, + identity_scaling_factor=input_tensor_scaling_factor, + ) + hidden_states, hidden_states_scaling_factor = self.LayerNorm(hidden_states, hidden_states_scaling_factor) + + hidden_states, hidden_states_scaling_factor = self.output_activation( + hidden_states, hidden_states_scaling_factor + ) + return hidden_states, hidden_states_scaling_factor + + +class IBertAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.quant_mode = config.quant_mode + self.self = IBertSelfAttention(config) + self.output = IBertSelfOutput(config) + + def forward( + self, + hidden_states, + hidden_states_scaling_factor, + attention_mask=None, + output_attentions=False, + ): + self_outputs, self_outputs_scaling_factor = self.self( + hidden_states, + hidden_states_scaling_factor, + attention_mask, + output_attentions, + ) + attention_output, attention_output_scaling_factor = self.output( + self_outputs[0], self_outputs_scaling_factor[0], hidden_states, hidden_states_scaling_factor + ) + outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them + outputs_scaling_factor = (attention_output_scaling_factor,) + self_outputs_scaling_factor[1:] + return outputs, outputs_scaling_factor + + +class IBertIntermediate(nn.Module): + def __init__(self, config): + super().__init__() + self.quant_mode = config.quant_mode + self.act_bit = 8 + self.weight_bit = 8 + self.bias_bit = 32 + self.dense = QuantLinear( + config.hidden_size, + config.intermediate_size, + bias=True, + weight_bit=self.weight_bit, + bias_bit=self.bias_bit, + quant_mode=self.quant_mode, + per_channel=True, + ) + if config.hidden_act != "gelu": + raise ValueError("I-BERT only supports 'gelu' for `config.hidden_act`") + self.intermediate_act_fn = IntGELU(quant_mode=self.quant_mode, force_dequant=config.force_dequant) + self.output_activation = QuantAct(self.act_bit, quant_mode=self.quant_mode) + + def forward(self, hidden_states, hidden_states_scaling_factor): + hidden_states, hidden_states_scaling_factor = self.dense(hidden_states, hidden_states_scaling_factor) + hidden_states, hidden_states_scaling_factor = self.intermediate_act_fn( + hidden_states, hidden_states_scaling_factor + ) + + # Requantization: 32bit -> 8-bit + hidden_states, hidden_states_scaling_factor = self.output_activation( + hidden_states, hidden_states_scaling_factor + ) + return hidden_states, hidden_states_scaling_factor + + +class IBertOutput(nn.Module): + def __init__(self, config): + super().__init__() + self.quant_mode = config.quant_mode + self.act_bit = 8 + self.weight_bit = 8 + self.bias_bit = 32 + self.ln_input_bit = 22 + self.ln_output_bit = 32 + + self.dense = QuantLinear( + config.intermediate_size, + config.hidden_size, + bias=True, + weight_bit=self.weight_bit, + bias_bit=self.bias_bit, + quant_mode=self.quant_mode, + per_channel=True, + ) + self.ln_input_act = QuantAct(self.ln_input_bit, quant_mode=self.quant_mode) + self.LayerNorm = IntLayerNorm( + config.hidden_size, + eps=config.layer_norm_eps, + output_bit=self.ln_output_bit, + quant_mode=self.quant_mode, + force_dequant=config.force_dequant, + ) + self.output_activation = QuantAct(self.act_bit, quant_mode=self.quant_mode) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward(self, hidden_states, hidden_states_scaling_factor, input_tensor, input_tensor_scaling_factor): + hidden_states, hidden_states_scaling_factor = self.dense(hidden_states, hidden_states_scaling_factor) + hidden_states = self.dropout(hidden_states) + hidden_states, hidden_states_scaling_factor = self.ln_input_act( + hidden_states, + hidden_states_scaling_factor, + identity=input_tensor, + identity_scaling_factor=input_tensor_scaling_factor, + ) + hidden_states, hidden_states_scaling_factor = self.LayerNorm(hidden_states, hidden_states_scaling_factor) + + hidden_states, hidden_states_scaling_factor = self.output_activation( + hidden_states, hidden_states_scaling_factor + ) + return hidden_states, hidden_states_scaling_factor + + +class IBertLayer(nn.Module): + def __init__(self, config): + super().__init__() + self.quant_mode = config.quant_mode + self.act_bit = 8 + + self.seq_len_dim = 1 + self.attention = IBertAttention(config) + self.intermediate = IBertIntermediate(config) + self.output = IBertOutput(config) + + self.pre_intermediate_act = QuantAct(self.act_bit, quant_mode=self.quant_mode) + self.pre_output_act = QuantAct(self.act_bit, quant_mode=self.quant_mode) + + def forward( + self, + hidden_states, + hidden_states_scaling_factor, + attention_mask=None, + output_attentions=False, + ): + self_attention_outputs, self_attention_outputs_scaling_factor = self.attention( + hidden_states, + hidden_states_scaling_factor, + attention_mask, + output_attentions=output_attentions, + ) + attention_output = self_attention_outputs[0] + attention_output_scaling_factor = self_attention_outputs_scaling_factor[0] + + outputs = self_attention_outputs[1:] # add self attentions if we output attention weights + + layer_output, layer_output_scaling_factor = self.feed_forward_chunk( + attention_output, attention_output_scaling_factor + ) + outputs = (layer_output,) + outputs + + return outputs + + def feed_forward_chunk(self, attention_output, attention_output_scaling_factor): + attention_output, attention_output_scaling_factor = self.pre_intermediate_act( + attention_output, attention_output_scaling_factor + ) + intermediate_output, intermediate_output_scaling_factor = self.intermediate( + attention_output, attention_output_scaling_factor + ) + + intermediate_output, intermediate_output_scaling_factor = self.pre_output_act( + intermediate_output, intermediate_output_scaling_factor + ) + layer_output, layer_output_scaling_factor = self.output( + intermediate_output, intermediate_output_scaling_factor, attention_output, attention_output_scaling_factor + ) + return layer_output, layer_output_scaling_factor + + +class IBertEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.quant_mode = config.quant_mode + self.layer = nn.ModuleList([IBertLayer(config) for _ in range(config.num_hidden_layers)]) + + def forward( + self, + hidden_states, + hidden_states_scaling_factor, + attention_mask=None, + output_attentions=False, + output_hidden_states=False, + return_dict=True, + ): + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + all_cross_attentions = None # `config.add_cross_attention` is not supported + + 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_scaling_factor, + 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, + all_cross_attentions, + ] + if v is not None + ) + return BaseModelOutputWithPastAndCrossAttentions( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + cross_attentions=all_cross_attentions, + ) + + +class IBertPooler(nn.Module): + def __init__(self, config): + super().__init__() + self.quant_mode = config.quant_mode + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.activation = nn.Tanh() + + def forward(self, hidden_states): + # 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 IBertPreTrainedModel(PreTrainedModel): + config: IBertConfig + base_model_prefix = "ibert" + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + if isinstance(module, (QuantLinear, nn.Linear)): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + init.zeros_(module.bias) + if getattr(module, "weight_integer", None) is not None: + init.zeros_(module.weight_integer) + init.zeros_(module.fc_scaling_factor) + if getattr(module, "bias_integer", None) is not None: + init.zeros_(module.bias_integer) + elif isinstance(module, (QuantEmbedding, nn.Embedding)): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + # 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]) + if getattr(module, "weight_scaling_factor", None) is not None: + init.zeros_(module.weight_scaling_factor) + init.zeros_(module.weight_integer) + elif isinstance(module, (IntLayerNorm, nn.LayerNorm)): + init.zeros_(module.bias) + init.ones_(module.weight) + if getattr(module, "shift", None) is not None: + init.zeros_(module.shift) + elif isinstance(module, IBertLMHead): + init.zeros_(module.bias) + elif isinstance(module, IBertEmbeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + elif isinstance(module, QuantAct): + init.constant_(module.x_min, -1e-5) + init.constant_(module.x_max, 1e-5) + init.zeros_(module.act_scaling_factor) + + def resize_token_embeddings(self, new_num_tokens=None): + raise NotImplementedError("`resize_token_embeddings` is not supported for I-BERT.") + + +@auto_docstring +class IBertModel(IBertPreTrainedModel): + """ + + 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. + + """ + + 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.quant_mode = config.quant_mode + + self.embeddings = IBertEmbeddings(config) + self.encoder = IBertEncoder(config) + + self.pooler = IBertPooler(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, + 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, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> BaseModelOutputWithPoolingAndCrossAttentions | tuple[torch.FloatTensor]: + 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") + + 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) + if token_type_ids is None: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, 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, embedding_output_scaling_factor = 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, + embedding_output_scaling_factor, + attention_mask=extended_attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = encoder_outputs[0] + 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 BaseModelOutputWithPoolingAndCrossAttentions( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + cross_attentions=encoder_outputs.cross_attentions, + ) + + +@auto_docstring +class IBertForMaskedLM(IBertPreTrainedModel): + _tied_weights_keys = { + "lm_head.decoder.weight": "ibert.embeddings.word_embeddings.weight$", + "lm_head.decoder.bias": "lm_head.bias", + } + + def __init__(self, config): + super().__init__(config) + + self.ibert = IBertModel(config, add_pooling_layer=False) + self.lm_head = IBertLMHead(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_output_embeddings(self): + return self.lm_head.decoder + + def set_output_embeddings(self, new_embeddings): + self.lm_head.decoder = new_embeddings + self.lm_head.bias = new_embeddings.bias + + @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, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> MaskedLMOutput | tuple[torch.FloatTensor]: + 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.ibert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = outputs[0] + prediction_scores = self.lm_head(sequence_output) + + masked_lm_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + 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, + attentions=outputs.attentions, + ) + + +class IBertLMHead(nn.Module): + """I-BERT Head for masked language modeling.""" + + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + self.decoder = nn.Linear(config.hidden_size, config.vocab_size) + self.bias = nn.Parameter(torch.zeros(config.vocab_size)) + + def forward(self, features, **kwargs): + x = self.dense(features) + x = gelu(x) + x = self.layer_norm(x) + + # project back to size of vocabulary with bias + x = self.decoder(x) + + return x + + +@auto_docstring( + custom_intro=""" + I-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 IBertForSequenceClassification(IBertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.ibert = IBertModel(config, add_pooling_layer=False) + self.classifier = IBertClassificationHead(config) + + # 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, + token_type_ids: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: 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, + ) -> SequenceClassifierOutput | tuple[torch.FloatTensor]: + 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.ibert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = outputs[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) + 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, + attentions=outputs.attentions, + ) + + +@auto_docstring +class IBertForMultipleChoice(IBertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.ibert = IBertModel(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.LongTensor | None = None, + token_type_ids: torch.LongTensor | None = None, + attention_mask: torch.FloatTensor | None = None, + labels: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.FloatTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> MultipleChoiceModelOutput | tuple[torch.FloatTensor]: + 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) + 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) + 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. + """ + 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] + + flat_input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None + flat_position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None + flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None + flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None + flat_inputs_embeds = ( + inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) + if inputs_embeds is not None + else None + ) + + outputs = self.ibert( + flat_input_ids, + position_ids=flat_position_ids, + token_type_ids=flat_token_type_ids, + attention_mask=flat_attention_mask, + inputs_embeds=flat_inputs_embeds, + output_attentions=output_attentions, + 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, + attentions=outputs.attentions, + ) + + +@auto_docstring +class IBertForTokenClassification(IBertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.ibert = IBertModel(config, add_pooling_layer=False) + 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.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, + labels: torch.LongTensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> TokenClassifierOutput | tuple[torch.FloatTensor]: + 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.ibert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + 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() + 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, + attentions=outputs.attentions, + ) + + +class IBertClassificationHead(nn.Module): + """Head for sentence-level classification tasks.""" + + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.out_proj = nn.Linear(config.hidden_size, config.num_labels) + + def forward(self, features, **kwargs): + hidden_states = features[:, 0, :] # take token (equiv. to [CLS]) + 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 IBertForQuestionAnswering(IBertPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.ibert = IBertModel(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() + + @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, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> QuestionAnsweringModelOutput | tuple[torch.FloatTensor]: + return_dict = return_dict if return_dict is not None else self.config.return_dict + + outputs = self.ibert( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + 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, + ) + + +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: + input_ids (`torch.LongTensor`): + Indices of input sequence tokens in the vocabulary. + + 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 + + +__all__ = [ + "IBertForMaskedLM", + "IBertForMultipleChoice", + "IBertForQuestionAnswering", + "IBertForSequenceClassification", + "IBertForTokenClassification", + "IBertModel", + "IBertPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/ibert/quant_modules.py b/third_party/transformers/src/transformers/models/ibert/quant_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..b3c6115bd10e08f6d5d23e53173dfe67caf92ed8 --- /dev/null +++ b/third_party/transformers/src/transformers/models/ibert/quant_modules.py @@ -0,0 +1,819 @@ +# Copyright 2021 The I-BERT Authors (Sehoon Kim, Amir Gholami, Zhewei Yao, +# Michael Mahoney, Kurt Keutzer - UC Berkeley) and The HuggingFace Inc. team. +# Copyright (c) 20121, 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. + +import decimal + +import numpy as np +import torch +from torch import nn +from torch.autograd import Function + +from ...utils import logging + + +logger = logging.get_logger(__name__) + + +class QuantEmbedding(nn.Module): + """ + Quantized version of `torch.nn.Embedding`. Adds quantization-specific arguments on top of `torch.nn.Embedding`. + + Args: + weight_bit (`int`, *optional*, defaults to `8`): + Bitwidth for the quantized weight. + momentum (`float`, *optional*, defaults to `0.95`): + Momentum for updating the activation quantization range. + quant_mode (`bool`, *optional*, defaults to `False`): + Whether or not the layer is quantized. + """ + + def __init__( + self, + num_embeddings, + embedding_dim, + padding_idx=None, + max_norm=None, + norm_type=2.0, + scale_grad_by_freq=False, + sparse=False, + _weight=None, + weight_bit=8, + momentum=0.95, + quant_mode=False, + ): + super().__init__() + self.num_ = num_embeddings + self.dim = embedding_dim + self.padding_idx = padding_idx + self.max_norm = max_norm + self.norm_type = norm_type + self.scale_grad_by_freq = scale_grad_by_freq + self.sparse = sparse + + self.weight = nn.Parameter(torch.zeros([num_embeddings, embedding_dim])) + self.register_buffer("weight_scaling_factor", torch.zeros(1)) + self.register_buffer("weight_integer", torch.zeros_like(self.weight)) + + self.weight_bit = weight_bit + self.momentum = momentum + self.quant_mode = quant_mode + self.percentile_mode = False + self.weight_function = SymmetricQuantFunction.apply + + def forward(self, x, positions=None, incremental_state=None): + if not self.quant_mode: + return ( + nn.functional.embedding( + x, + self.weight, + self.padding_idx, + self.max_norm, + self.norm_type, + self.scale_grad_by_freq, + self.sparse, + ), + None, + ) + + w = self.weight + w_transform = w.data.detach() + w_min = w_transform.min().expand(1) + w_max = w_transform.max().expand(1) + + self.weight_scaling_factor = symmetric_linear_quantization_params(self.weight_bit, w_min, w_max, False) + self.weight_integer = self.weight_function( + self.weight, self.weight_bit, self.percentile_mode, self.weight_scaling_factor + ) + + emb_int = nn.functional.embedding( + x, + self.weight_integer, + self.padding_idx, + self.max_norm, + self.norm_type, + self.scale_grad_by_freq, + self.sparse, + ) + return emb_int * self.weight_scaling_factor, self.weight_scaling_factor + + +class QuantAct(nn.Module): + """ + Quantizes the given activation. + + Args: + activation_bit (`int`): + Bitwidth for the quantized activation. + act_range_momentum (`float`, *optional*, defaults to `0.95`): + Momentum for updating the activation quantization range. + per_channel (`bool`, *optional*, defaults to `False`): + Whether to or not use channel-wise quantization. + channel_len (`int`, *optional*): + Specify the channel length when set the *per_channel* True. + quant_mode (`bool`, *optional*, defaults to `False`): + Whether or not the layer is quantized. + """ + + def __init__(self, activation_bit, act_range_momentum=0.95, per_channel=False, channel_len=None, quant_mode=False): + super().__init__() + + self.activation_bit = activation_bit + self.act_range_momentum = act_range_momentum + self.quant_mode = quant_mode + self.per_channel = per_channel + self.percentile = False + self.act_function = SymmetricQuantFunction.apply + + if not self.per_channel: + self.register_buffer("x_min", torch.zeros(1)) + self.register_buffer("x_max", torch.zeros(1)) + self.register_buffer("act_scaling_factor", torch.zeros(1)) + self.x_min -= 1e-5 + self.x_max += 1e-5 + else: + raise NotImplementedError("per-channel mode is not currently supported for activation.") + + def __repr__(self): + return ( + f"{self.__class__.__name__}(activation_bit={self.activation_bit}, " + f"quant_mode: {self.quant_mode}, Act_min: {self.x_min.item():.2f}, " + f"Act_max: {self.x_max.item():.2f})" + ) + + def forward( + self, + x, + pre_act_scaling_factor=None, + identity=None, + identity_scaling_factor=None, + specified_min=None, + specified_max=None, + ): + x_act = x if identity is None else identity + x + # collect running stats if training + if self.training: + assert not self.percentile, "percentile mode is not currently supported for activation." + assert not self.per_channel, "per-channel mode is not currently supported for activation." + x_min = x_act.data.min() + x_max = x_act.data.max() + + assert x_max.isnan().sum() == 0 and x_min.isnan().sum() == 0, ( + "NaN detected when computing min/max of the activation" + ) + + # Initialization + if self.x_min.min() > -1.1e-5 and self.x_max.max() < 1.1e-5: + self.x_min = self.x_min + x_min + self.x_max = self.x_max + x_max + + # exponential moving average (EMA) + # use momentum to prevent the quantized values change greatly every iteration + elif self.act_range_momentum == -1: + self.x_min = torch.min(self.x_min, x_min) + self.x_max = torch.max(self.x_max, x_max) + else: + self.x_min = self.x_min * self.act_range_momentum + x_min * (1 - self.act_range_momentum) + self.x_max = self.x_max * self.act_range_momentum + x_max * (1 - self.act_range_momentum) + + if not self.quant_mode: + return x_act, None + + x_min = self.x_min if specified_min is None else specified_min + x_max = self.x_max if specified_max is None else specified_max + + self.act_scaling_factor = symmetric_linear_quantization_params( + self.activation_bit, x_min, x_max, per_channel=self.per_channel + ) + + if pre_act_scaling_factor is None: + # this is for the input quantization + quant_act_int = self.act_function(x, self.activation_bit, self.percentile, self.act_scaling_factor) + else: + quant_act_int = FixedPointMul.apply( + x, + pre_act_scaling_factor, + self.activation_bit, + self.act_scaling_factor, + identity, + identity_scaling_factor, + ) + + correct_output_scale = self.act_scaling_factor.view(-1) + + return quant_act_int * correct_output_scale, self.act_scaling_factor + + +class QuantLinear(nn.Module): + """ + Quantized version of `torch.nn.Linear`. Adds quantization-specific arguments on top of `torch.nn.Linear`. + + Args: + weight_bit (`int`, *optional*, defaults to `8`): + Bitwidth for the quantized weight. + bias_bit (`int`, *optional*, defaults to `32`): + Bitwidth for the quantized bias. + per_channel (`bool`, *optional*, defaults to `False`): + Whether or not to use channel-wise quantization. + quant_mode (`bool`, *optional*, defaults to `False`): + Whether or not the layer is quantized. + """ + + def __init__( + self, in_features, out_features, bias=True, weight_bit=8, bias_bit=32, per_channel=False, quant_mode=False + ): + super().__init__() + self.in_features = in_features + self.out_features = out_features + + self.weight = nn.Parameter(torch.zeros([out_features, in_features])) + self.register_buffer("weight_integer", torch.zeros_like(self.weight)) + self.register_buffer("fc_scaling_factor", torch.zeros(self.out_features)) + if bias: + self.bias = nn.Parameter(torch.zeros(out_features)) + self.register_buffer("bias_integer", torch.zeros_like(self.bias)) + + self.weight_bit = weight_bit + self.quant_mode = quant_mode + self.per_channel = per_channel + self.bias_bit = bias_bit + self.quant_mode = quant_mode + self.percentile_mode = False + self.weight_function = SymmetricQuantFunction.apply + + def __repr__(self): + s = super().__repr__() + s = f"({s} weight_bit={self.weight_bit}, quant_mode={self.quant_mode})" + return s + + def forward(self, x, prev_act_scaling_factor=None): + if not self.quant_mode: + return nn.functional.linear(x, weight=self.weight, bias=self.bias), None + + # assert that prev_act_scaling_factor is a scalar tensor + assert prev_act_scaling_factor is not None and prev_act_scaling_factor.shape == (1,), ( + "Input activation to the QuantLinear layer should be globally (non-channel-wise) quantized. " + "Please add a QuantAct layer with `per_channel = True` before this QuantAct layer" + ) + + w = self.weight + w_transform = w.data.detach() + if self.per_channel: + w_min, _ = torch.min(w_transform, dim=1, out=None) + w_max, _ = torch.max(w_transform, dim=1, out=None) + else: + w_min = w_transform.min().expand(1) + w_max = w_transform.max().expand(1) + + self.fc_scaling_factor = symmetric_linear_quantization_params(self.weight_bit, w_min, w_max, self.per_channel) + self.weight_integer = self.weight_function( + self.weight, self.weight_bit, self.percentile_mode, self.fc_scaling_factor + ) + + bias_scaling_factor = self.fc_scaling_factor * prev_act_scaling_factor + + if self.bias is not None: + self.bias_integer = self.weight_function(self.bias, self.bias_bit, False, bias_scaling_factor) + + prev_act_scaling_factor = prev_act_scaling_factor.view(1, -1) + x_int = x / prev_act_scaling_factor + + return ( + nn.functional.linear(x_int, weight=self.weight_integer, bias=self.bias_integer) * bias_scaling_factor, + bias_scaling_factor, + ) + + +class IntGELU(nn.Module): + """ + Quantized version of `torch.nn.GELU`. Adds quantization-specific arguments on top of `torch.nn.GELU`. + + Args: + quant_mode (`bool`, *optional*, defaults to `False`): + Whether or not the layer is quantized. + force_dequant (`str`, *optional*, defaults to `"none"`): + Force dequantize the layer if either "gelu" or "nonlinear" is given. + """ + + def __init__(self, quant_mode=True, force_dequant="none"): + super().__init__() + self.quant_mode = quant_mode + + if force_dequant in ["nonlinear", "gelu"]: + logger.info("Force dequantize gelu") + self.quant_mode = False + + if not self.quant_mode: + self.activation_fn = nn.GELU() + + self.k = 1.4142 + self.const = 14 # dummy integer constant + self.coeff = [-0.2888, -1.769, 1] # a(x+b)**2 + c + self.coeff[2] /= self.coeff[0] + + def int_erf(self, x_int, scaling_factor): + b_int = torch.floor(self.coeff[1] / scaling_factor) + c_int = torch.floor(self.coeff[2] / scaling_factor**2) + sign = torch.sign(x_int) + + abs_int = torch.min(torch.abs(x_int), -b_int) + y_int = sign * ((abs_int + b_int) ** 2 + c_int) + scaling_factor = scaling_factor**2 * self.coeff[0] + + # avoid overflow + y_int = floor_ste.apply(y_int / 2**self.const) + scaling_factor = scaling_factor * 2**self.const + + return y_int, scaling_factor + + def forward(self, x, scaling_factor=None): + if not self.quant_mode: + return self.activation_fn(x), None + + x_int = x / scaling_factor + sigmoid_int, sigmoid_scaling_factor = self.int_erf(x_int, scaling_factor / self.k) + + shift_int = 1.0 // sigmoid_scaling_factor + + x_int = x_int * (sigmoid_int + shift_int) + scaling_factor = scaling_factor * sigmoid_scaling_factor / 2 + + return x_int * scaling_factor, scaling_factor + + +class IntSoftmax(nn.Module): + """ + Quantized version of `torch.nn.Softmax`. Adds quantization-specific arguments on top of `torch.nn.Softmax`. + + Args: + output_bit (`int`): + Bitwidth for the layer output activation. + quant_mode (`bool`, *optional*, defaults to `False`): + Whether or not the layer is quantized. + force_dequant (`str`, *optional*, defaults to `"none"`): + Force dequantize the layer if either "softmax" or "nonlinear" is given. + """ + + def __init__(self, output_bit, quant_mode=False, force_dequant="none"): + super().__init__() + self.output_bit = output_bit + self.max_bit = 32 + self.quant_mode = quant_mode + + if force_dequant in ["nonlinear", "softmax"]: + logger.info("Force dequantize softmax") + self.quant_mode = False + + self.act = QuantAct(16, quant_mode=self.quant_mode) + self.x0 = -0.6931 # -ln2 + self.const = 30 # dummy integer constant + self.coef = [0.35815147, 0.96963238, 1.0] # ax**2 + bx + c + self.coef[1] /= self.coef[0] + self.coef[2] /= self.coef[0] + + def int_polynomial(self, x_int, scaling_factor): + with torch.no_grad(): + b_int = torch.floor(self.coef[1] / scaling_factor) + c_int = torch.floor(self.coef[2] / scaling_factor**2) + z = (x_int + b_int) * x_int + c_int + scaling_factor = self.coef[0] * scaling_factor**2 + return z, scaling_factor + + def int_exp(self, x_int, scaling_factor): + with torch.no_grad(): + x0_int = torch.floor(self.x0 / scaling_factor) + x_int = torch.max(x_int, self.const * x0_int) + + q = floor_ste.apply(x_int / x0_int) + r = x_int - x0_int * q + exp_int, exp_scaling_factor = self.int_polynomial(r, scaling_factor) + exp_int = torch.clamp(floor_ste.apply(exp_int * 2 ** (self.const - q)), min=0) + scaling_factor = exp_scaling_factor / 2**self.const + return exp_int, scaling_factor + + def forward(self, x, scaling_factor): + if not self.quant_mode: + return nn.functional.softmax(x, dim=-1), None + + x_int = x / scaling_factor + + x_int_max, _ = x_int.max(dim=-1, keepdim=True) + x_int = x_int - x_int_max + exp_int, exp_scaling_factor = self.int_exp(x_int, scaling_factor) + + # Avoid overflow + exp, exp_scaling_factor = self.act(exp_int, exp_scaling_factor) + exp_int = exp / exp_scaling_factor + + exp_int_sum = exp_int.sum(dim=-1, keepdim=True) + factor = floor_ste.apply(2**self.max_bit / exp_int_sum) + exp_int = floor_ste.apply(exp_int * factor / 2 ** (self.max_bit - self.output_bit)) + scaling_factor = 1 / 2**self.output_bit + return exp_int * scaling_factor, scaling_factor + + +class IntLayerNorm(nn.Module): + """ + Quantized version of `torch.nn.LayerNorm`. Adds quantization-specific arguments on top of `torch.nn.LayerNorm`. + + Args: + output_bit (`int`, *optional*, defaults to `8`): + Bitwidth for the layer output activation. + quant_mode (`bool`, *optional*, defaults to `False`): + Whether or not the layer is quantized. + force_dequant (`str`, *optional*, defaults to `"none"`): + Force dequantize the layer if either "layernorm" or "nonlinear" is given. + """ + + def __init__(self, normalized_shape, eps, output_bit=8, quant_mode=False, force_dequant="none"): + super().__init__() + self.normalized_shape = normalized_shape + self.eps = eps + + self.weight = nn.Parameter(torch.zeros(normalized_shape)) + self.bias = nn.Parameter(torch.zeros(normalized_shape)) + + self.quant_mode = quant_mode + if force_dequant in ["nonlinear", "layernorm"]: + logger.info("Force dequantize layernorm") + self.quant_mode = False + + self.register_buffer("shift", torch.zeros(1)) + self.output_bit = output_bit + self.max_bit = 32 + self.dim_sqrt = None + self.activation = QuantAct(self.output_bit, quant_mode=self.quant_mode) + + def set_shift(self, y_int): + with torch.no_grad(): + y_sq_int = y_int**2 + var_int = torch.sum(y_sq_int, axis=2, keepdim=True) + shift = (torch.log2(torch.sqrt(var_int / 2**self.max_bit)).ceil()).max() + shift_old = self.shift + self.shift = torch.max(self.shift, shift) + logger.info(f"Dynamic shift adjustment: {int(shift_old)} -> {int(self.shift)}") + + def overflow_fallback(self, y_int): + """ + This fallback function is called when overflow is detected during training time, and adjusts the `self.shift` + to avoid overflow in the subsequent runs. + """ + self.set_shift(y_int) # adjusts `self.shift` + y_int_shifted = floor_ste.apply(y_int / 2**self.shift) + y_sq_int = y_int_shifted**2 + var_int = torch.sum(y_sq_int, axis=2, keepdim=True) + return var_int + + def forward(self, x, scaling_factor=None): + if not self.quant_mode: + mean = x.mean(axis=2, keepdim=True) + y = x - mean + var = torch.mean(y**2, axis=2, keepdim=True) + x = y / torch.sqrt(self.eps + var) + x = x * self.weight + self.bias + return x, None + + # compute sqrt of the feature dimension if it is the first run + if self.dim_sqrt is None: + n = torch.tensor(x.shape[2], dtype=torch.float) + self.dim_sqrt = torch.sqrt(n).to(x.device) + + # Normalization: computes mean and variance(std) + x_int = x / scaling_factor + mean_int = round_ste.apply(x_int.mean(axis=2, keepdim=True)) + y_int = x_int - mean_int + y_int_shifted = floor_ste.apply(y_int / 2**self.shift) + y_sq_int = y_int_shifted**2 + var_int = torch.sum(y_sq_int, axis=2, keepdim=True) + + # overflow handling in training time + if self.training: + # if overflow is detected + if var_int.max() >= 2**self.max_bit: + var_int = self.overflow_fallback(y_int) + assert var_int.max() < 2**self.max_bit + 0.1, ( + "Error detected in overflow handling: " + "`var_int` exceeds `self.max_bit` (the maximum possible bit width)" + ) + + # To be replaced with integer-sqrt kernel that produces the same output + std_int = floor_ste.apply(torch.sqrt(var_int)) * 2**self.shift + factor = floor_ste.apply(2**31 / std_int) + y_int = floor_ste.apply(y_int * factor / 2) + scaling_factor = self.dim_sqrt / 2**30 + + # scaling and shifting + bias = self.bias.data.detach() / (self.weight.data.detach()) + bias_int = floor_ste.apply(bias / scaling_factor) + + y_int = y_int + bias_int + scaling_factor = scaling_factor * self.weight + x = y_int * scaling_factor + + return x, scaling_factor + + +def get_percentile_min_max(input, lower_percentile, upper_percentile, output_tensor=False): + """ + Calculate the percentile max and min values in a given tensor + + Args: + input (`torch.Tensor`): + The target tensor to calculate percentile max and min. + lower_percentile (`float`): + If 0.1, means we return the value of the smallest 0.1% value in the tensor as percentile min. + upper_percentile (`float`): + If 99.9, means we return the value of the largest 0.1% value in the tensor as percentile max. + output_tensor (`bool`, *optional*, defaults to `False`): + If True, this function returns tensors, otherwise it returns values. + + Returns: + `Tuple(torch.Tensor, torch.Tensor)`: Percentile min and max value of *input* + """ + input_length = input.shape[0] + + lower_index = round(input_length * (1 - lower_percentile * 0.01)) + upper_index = round(input_length * upper_percentile * 0.01) + + upper_bound = torch.kthvalue(input, k=upper_index).values + + if lower_percentile == 0: + lower_bound = upper_bound * 0 + # lower_index += 1 + else: + lower_bound = -torch.kthvalue(-input, k=lower_index).values + + if not output_tensor: + lower_bound = lower_bound.item() + upper_bound = upper_bound.item() + return lower_bound, upper_bound + + +def linear_quantize(input, scale, zero_point, inplace=False): + """ + Quantize single-precision input tensor to integers with the given scaling factor and zeropoint. + + Args: + input (`torch.Tensor`): + Single-precision input tensor to be quantized. + scale (`torch.Tensor`): + Scaling factor for quantization. + zero_pint (`torch.Tensor`): + Shift for quantization. + inplace (`bool`, *optional*, defaults to `False`): + Whether to compute inplace or not. + + Returns: + `torch.Tensor`: Linearly quantized value of *input* according to *scale* and *zero_point*. + """ + # reshape scale and zeropoint for convolutional weights and activation + if len(input.shape) == 4: + scale = scale.view(-1, 1, 1, 1) + zero_point = zero_point.view(-1, 1, 1, 1) + # reshape scale and zeropoint for linear weights + elif len(input.shape) == 2: + scale = scale.view(-1, 1) + zero_point = zero_point.view(-1, 1) + else: + scale = scale.view(-1) + zero_point = zero_point.view(-1) + # quantized = float / scale + zero_point + if inplace: + input.mul_(1.0 / scale).add_(zero_point).round_() + return input + return torch.round(1.0 / scale * input + zero_point) + + +def symmetric_linear_quantization_params(num_bits, saturation_min, saturation_max, per_channel=False): + """ + Compute the scaling factor with the given quantization range for symmetric quantization. + + Args: + saturation_min (`torch.Tensor`): + Lower bound for quantization range. + saturation_max (`torch.Tensor`): + Upper bound for quantization range. + per_channel (`bool`, *optional*, defaults to `False`): + Whether to or not use channel-wise quantization. + + Returns: + `torch.Tensor`: Scaling factor that linearly quantizes the given range between *saturation_min* and + *saturation_max*. + """ + # in this part, we do not need any gradient computation, + # in order to enforce this, we put torch.no_grad() + with torch.no_grad(): + n = 2 ** (num_bits - 1) - 1 + + if per_channel: + scale, _ = torch.max(torch.stack([saturation_min.abs(), saturation_max.abs()], dim=1), dim=1) + scale = torch.clamp(scale, min=1e-8) / n + + else: + scale = max(saturation_min.abs(), saturation_max.abs()) + scale = torch.clamp(scale, min=1e-8) / n + + return scale + + +class SymmetricQuantFunction(Function): + """ + Class to quantize the given floating-point values using symmetric quantization with given range and bitwidth. + """ + + @staticmethod + def forward(ctx, x, k, percentile_mode, scale): + """ + Args: + x (`torch.Tensor`): + Floating point tensor to be quantized. + k (`int`): + Quantization bitwidth. + percentile_mode (`bool`): + Whether or not to use percentile calibration. + scale (`torch.Tensor`): + Pre-calculated scaling factor for *x*. Note that the current implementation of SymmetricQuantFunction + requires pre-calculated scaling factor. + + Returns: + `torch.Tensor`: Symmetric-quantized value of *input*. + """ + zero_point = torch.tensor(0.0, device=scale.device) + + n = 2 ** (k - 1) - 1 + new_quant_x = linear_quantize(x, scale, zero_point, inplace=False) + new_quant_x = torch.clamp(new_quant_x, -n, n - 1) + + ctx.scale = scale + return new_quant_x + + @staticmethod + def backward(ctx, grad_output): + scale = ctx.scale + if len(grad_output.shape) == 4: + scale = scale.view(-1, 1, 1, 1) + # reshape scale and zeropoint for linear weights + elif len(grad_output.shape) == 2: + scale = scale.view(-1, 1) + else: + scale = scale.view(-1) + + return grad_output.clone() / scale, None, None, None, None + + +class floor_ste(Function): + """ + Straight-through Estimator(STE) for torch.floor() + """ + + @staticmethod + def forward(ctx, x): + return torch.floor(x) + + @staticmethod + def backward(ctx, grad_output): + return grad_output.clone() + + +class round_ste(Function): + """ + Straight-through Estimator(STE) for torch.round() + """ + + @staticmethod + def forward(ctx, x): + return torch.round(x) + + @staticmethod + def backward(ctx, grad_output): + return grad_output.clone() + + +def batch_frexp(inputs, max_bit=31): + """ + Decompose the scaling factor into mantissa and twos exponent. + + Args: + scaling_factor (`torch.Tensor`): + Target scaling factor to decompose. + + Returns: + ``Tuple(torch.Tensor, torch.Tensor)`: mantisa and exponent + """ + + shape_of_input = inputs.size() + + # trans the input to be a 1-d tensor + inputs = inputs.view(-1) + + output_m, output_e = np.frexp(inputs.cpu().numpy()) + tmp_m = [] + for m in output_m: + int_m_shifted = int( + decimal.Decimal(m * (2**max_bit)).quantize(decimal.Decimal(1), rounding=decimal.ROUND_HALF_UP) + ) + tmp_m.append(int_m_shifted) + output_m = np.array(tmp_m) + + output_e = float(max_bit) - output_e + + return ( + torch.from_numpy(output_m).to(inputs.device).view(shape_of_input), + torch.from_numpy(output_e).to(inputs.device).view(shape_of_input), + ) + + +class FixedPointMul(Function): + """ + Function to perform fixed-point arithmetic that can match integer arithmetic on hardware. + + Args: + pre_act (`torch.Tensor`): + Input tensor. + pre_act_scaling_factor (`torch.Tensor`): + Scaling factor of the input tensor *pre_act*. + bit_num (`int`): + Quantization bitwidth. + z_scaling_factor (`torch.Tensor`): + Scaling factor of the output tensor. + identity (`torch.Tensor`, *optional*): + Identity tensor, if exists. + identity_scaling_factor (`torch.Tensor`, *optional*): + Scaling factor of the identity tensor *identity*, if exists. + + Returns: + `torch.Tensor`: Output tensor(*pre_act* if *identity* is not given, otherwise the addition of *pre_act* and + *identity*), whose scale is rescaled to *z_scaling_factor*. + """ + + @staticmethod + def forward( + ctx, + pre_act, + pre_act_scaling_factor, + bit_num, + z_scaling_factor, + identity=None, + identity_scaling_factor=None, + ): + if len(pre_act_scaling_factor.shape) == 3: + reshape = lambda x: x # noqa: E731 + else: + reshape = lambda x: x.view(1, 1, -1) # noqa: E731 + ctx.identity = identity + + n = 2 ** (bit_num - 1) - 1 + + with torch.no_grad(): + pre_act_scaling_factor = reshape(pre_act_scaling_factor) + if identity is not None: + identity_scaling_factor = reshape(identity_scaling_factor) + + ctx.z_scaling_factor = z_scaling_factor + + z_int = torch.round(pre_act / pre_act_scaling_factor) + _A = pre_act_scaling_factor.type(torch.double) + _B = (z_scaling_factor.type(torch.float)).type(torch.double) + new_scale = _A / _B + new_scale = reshape(new_scale) + + m, e = batch_frexp(new_scale) + + output = z_int.type(torch.double) * m.type(torch.double) + output = torch.round(output / (2.0**e)) + + if identity is not None: + # needs addition of identity activation + wx_int = torch.round(identity / identity_scaling_factor) + + _A = identity_scaling_factor.type(torch.double) + _B = (z_scaling_factor.type(torch.float)).type(torch.double) + new_scale = _A / _B + new_scale = reshape(new_scale) + + m1, e1 = batch_frexp(new_scale) + output1 = wx_int.type(torch.double) * m1.type(torch.double) + output1 = torch.round(output1 / (2.0**e1)) + + output = output1 + output + + return torch.clamp(output.type(torch.float), -n - 1, n) + + @staticmethod + def backward(ctx, grad_output): + identity_grad = None + if ctx.identity is not None: + identity_grad = grad_output.clone() / ctx.z_scaling_factor + return grad_output.clone() / ctx.z_scaling_factor, None, None, None, None, identity_grad, None diff --git a/third_party/transformers/src/transformers/models/layoutlm/__init__.py b/third_party/transformers/src/transformers/models/layoutlm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5edb5482146e8f8bee62d44ab1948446edb0f0de --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlm/__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 ..bert.tokenization_bert import BertTokenizer as LayoutLMTokenizer + from ..bert.tokenization_bert import BertTokenizer as LayoutLMTokenizerFast + from .configuration_layoutlm import * + from .modeling_layoutlm 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/layoutlm/configuration_layoutlm.py b/third_party/transformers/src/transformers/models/layoutlm/configuration_layoutlm.py new file mode 100644 index 0000000000000000000000000000000000000000..f76d15dcb00a57cdb7f9351f30eda68571bc0e58 --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlm/configuration_layoutlm.py @@ -0,0 +1,67 @@ +# Copyright 2010, The Microsoft Research Asia LayoutLM Team authors +# +# 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. +"""LayoutLM model configuration""" + +from huggingface_hub.dataclasses import strict + +from ... import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="microsoft/layoutlm-base-uncased") +@strict +class LayoutLMConfig(PreTrainedConfig): + r""" + max_2d_position_embeddings (`int`, *optional*, defaults to 1024): + The maximum value that the 2D position embedding might ever used. Typically set this to something large + just in case (e.g., 1024). + + Examples: + + ```python + >>> from transformers import LayoutLMConfig, LayoutLMModel + + >>> # Initializing a LayoutLM configuration + >>> configuration = LayoutLMConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = LayoutLMModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "layoutlm" + + 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 + eos_token_id: int | list[int] | None = None + bos_token_id: int | None = None + use_cache: bool = True + max_2d_position_embeddings: int = 1024 + tie_word_embeddings: bool = True + + +__all__ = ["LayoutLMConfig"] diff --git a/third_party/transformers/src/transformers/models/layoutlm/modeling_layoutlm.py b/third_party/transformers/src/transformers/models/layoutlm/modeling_layoutlm.py new file mode 100644 index 0000000000000000000000000000000000000000..65ebb99a62eb4fcb7b6398cfdc7521fa04055e62 --- /dev/null +++ b/third_party/transformers/src/transformers/models/layoutlm/modeling_layoutlm.py @@ -0,0 +1,1012 @@ +# Copyright 2018 The Microsoft Research Asia LayoutLM 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 LayoutLM model.""" + +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 ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutput, + BaseModelOutputWithPooling, + MaskedLMOutput, + 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 TransformersKwargs, auto_docstring, can_return_tuple, logging +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from .configuration_layoutlm import LayoutLMConfig + + +logger = logging.get_logger(__name__) + + +LayoutLMLayerNorm = nn.LayerNorm + + +class LayoutLMEmbeddings(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.x_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.hidden_size) + self.y_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.hidden_size) + self.h_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.hidden_size) + self.w_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.hidden_size) + self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size) + + self.LayerNorm = LayoutLMLayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + self.register_buffer( + "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False + ) + + def forward( + self, + input_ids=None, + bbox=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] + + device = input_ids.device if input_ids is not None else inputs_embeds.device + + if position_ids is None: + position_ids = self.position_ids[:, :seq_length] + + if token_type_ids is None: + token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device) + + if inputs_embeds is None: + inputs_embeds = self.word_embeddings(input_ids) + + words_embeddings = inputs_embeds + position_embeddings = self.position_embeddings(position_ids) + try: + left_position_embeddings = self.x_position_embeddings(bbox[:, :, 0]) + upper_position_embeddings = self.y_position_embeddings(bbox[:, :, 1]) + right_position_embeddings = self.x_position_embeddings(bbox[:, :, 2]) + lower_position_embeddings = self.y_position_embeddings(bbox[:, :, 3]) + except IndexError as e: + raise IndexError("The `bbox`coordinate values should be within 0-1000 range.") from e + + h_position_embeddings = self.h_position_embeddings(bbox[:, :, 3] - bbox[:, :, 1]) + w_position_embeddings = self.w_position_embeddings(bbox[:, :, 2] - bbox[:, :, 0]) + token_type_embeddings = self.token_type_embeddings(token_type_ids) + + embeddings = ( + words_embeddings + + position_embeddings + + left_position_embeddings + + upper_position_embeddings + + right_position_embeddings + + lower_position_embeddings + + h_position_embeddings + + w_position_embeddings + + token_type_embeddings + ) + embeddings = self.LayerNorm(embeddings) + embeddings = self.dropout(embeddings) + return embeddings + + +# 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 AlignText->LayoutLM +class LayoutLMSelfAttention(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 with Bert->LayoutLM +class LayoutLMSelfOutput(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 AlignText->LayoutLM +class LayoutLMAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.self = LayoutLMSelfAttention(config) + self.output = LayoutLMSelfOutput(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 LayoutLMIntermediate(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->LayoutLM +class LayoutLMOutput(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->LayoutLM +class LayoutLMLayer(GradientCheckpointingLayer): + def __init__(self, config): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = LayoutLMAttention(config) + self.intermediate = LayoutLMIntermediate(config) + self.output = LayoutLMOutput(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->LayoutLM +class LayoutLMEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.layer = nn.ModuleList([LayoutLMLayer(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 LayoutLMPooler(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->LayoutLM +class LayoutLMPredictionHeadTransform(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 + + +# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead with Bert->LayoutLM +class LayoutLMLMPredictionHead(nn.Module): + def __init__(self, config): + super().__init__() + self.transform = LayoutLMPredictionHeadTransform(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 + + +# Copied from transformers.models.bert.modeling_bert.BertOnlyMLMHead with Bert->LayoutLM +class LayoutLMOnlyMLMHead(nn.Module): + def __init__(self, config): + super().__init__() + self.predictions = LayoutLMLMPredictionHead(config) + + def forward(self, sequence_output: torch.Tensor) -> torch.Tensor: + prediction_scores = self.predictions(sequence_output) + return prediction_scores + + +@auto_docstring +class LayoutLMPreTrainedModel(PreTrainedModel): + config: LayoutLMConfig + base_model_prefix = "layoutlm" + supports_gradient_checkpointing = True + _can_record_outputs = { + "hidden_states": LayoutLMLayer, + "attentions": LayoutLMSelfAttention, + } + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights""" + super()._init_weights(module) + if isinstance(module, LayoutLMLMPredictionHead): + init.zeros_(module.bias) + elif isinstance(module, LayoutLMEmbeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1))) + + +@auto_docstring +class LayoutLMModel(LayoutLMPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.config = config + + self.embeddings = LayoutLMEmbeddings(config) + self.encoder = LayoutLMEncoder(config) + self.pooler = LayoutLMPooler(config) + + # 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.LongTensor | None = None, + bbox: 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, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + bbox (`torch.LongTensor` of shape `(batch_size, sequence_length, 4)`, *optional*): + Bounding boxes of each input sequence tokens. Selected in the range `[0, + config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1) + format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1, + y1) represents the position of the lower right corner. See [Overview](#Overview) for normalization. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, LayoutLMModel + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased") + >>> model = LayoutLMModel.from_pretrained("microsoft/layoutlm-base-uncased") + + >>> words = ["Hello", "world"] + >>> normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782] + + >>> token_boxes = [] + >>> for word, box in zip(words, normalized_word_boxes): + ... word_tokens = tokenizer.tokenize(word) + ... token_boxes.extend([box] * len(word_tokens)) + >>> # add bounding boxes of cls + sep tokens + >>> token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]] + + >>> encoding = tokenizer(" ".join(words), return_tensors="pt") + >>> input_ids = encoding["input_ids"] + >>> attention_mask = encoding["attention_mask"] + >>> token_type_ids = encoding["token_type_ids"] + >>> bbox = torch.tensor([token_boxes]) + + >>> outputs = model( + ... input_ids=input_ids, bbox=bbox, attention_mask=attention_mask, token_type_ids=token_type_ids + ... ) + + >>> last_hidden_states = outputs.last_hidden_state + ```""" + 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) + + if bbox is None: + bbox = torch.zeros(input_shape + (4,), dtype=torch.long, device=device) + + extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) + + extended_attention_mask = extended_attention_mask.to(dtype=self.dtype) + extended_attention_mask = (1.0 - extended_attention_mask) * torch.finfo(self.dtype).min + + embedding_output = self.embeddings( + input_ids=input_ids, + bbox=bbox, + position_ids=position_ids, + token_type_ids=token_type_ids, + inputs_embeds=inputs_embeds, + ) + encoder_outputs = self.encoder( + embedding_output, + extended_attention_mask, + **kwargs, + ) + sequence_output = encoder_outputs[0] + pooled_output = self.pooler(sequence_output) + + return BaseModelOutputWithPooling( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + ) + + +@auto_docstring +class LayoutLMForMaskedLM(LayoutLMPreTrainedModel): + _tied_weights_keys = { + "cls.predictions.decoder.bias": "cls.predictions.bias", + "cls.predictions.decoder.weight": "layoutlm.embeddings.word_embeddings.weight", + } + + def __init__(self, config): + super().__init__(config) + + self.layoutlm = LayoutLMModel(config) + self.cls = LayoutLMOnlyMLMHead(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.layoutlm.embeddings.word_embeddings + + 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.LongTensor | None = None, + bbox: 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, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | MaskedLMOutput: + r""" + bbox (`torch.LongTensor` of shape `(batch_size, sequence_length, 4)`, *optional*): + Bounding boxes of each input sequence tokens. Selected in the range `[0, + config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1) + format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1, + y1) represents the position of the lower right corner. See [Overview](#Overview) for normalization. + 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 AutoTokenizer, LayoutLMForMaskedLM + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased") + >>> model = LayoutLMForMaskedLM.from_pretrained("microsoft/layoutlm-base-uncased") + + >>> words = ["Hello", "[MASK]"] + >>> normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782] + + >>> token_boxes = [] + >>> for word, box in zip(words, normalized_word_boxes): + ... word_tokens = tokenizer.tokenize(word) + ... token_boxes.extend([box] * len(word_tokens)) + >>> # add bounding boxes of cls + sep tokens + >>> token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]] + + >>> encoding = tokenizer(" ".join(words), return_tensors="pt") + >>> input_ids = encoding["input_ids"] + >>> attention_mask = encoding["attention_mask"] + >>> token_type_ids = encoding["token_type_ids"] + >>> bbox = torch.tensor([token_boxes]) + + >>> labels = tokenizer("Hello world", return_tensors="pt")["input_ids"] + + >>> outputs = model( + ... input_ids=input_ids, + ... bbox=bbox, + ... attention_mask=attention_mask, + ... token_type_ids=token_type_ids, + ... labels=labels, + ... ) + + >>> loss = outputs.loss + ```""" + outputs = self.layoutlm( + input_ids, + bbox, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + sequence_output = outputs[0] + prediction_scores = self.cls(sequence_output) + + masked_lm_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + 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=""" + LayoutLM Model with a sequence classification head on top (a linear layer on top of the pooled output) e.g. for + document image classification tasks such as the [RVL-CDIP](https://www.cs.cmu.edu/~aharley/rvl-cdip/) dataset. + """ +) +class LayoutLMForSequenceClassification(LayoutLMPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.layoutlm = LayoutLMModel(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() + + def get_input_embeddings(self): + return self.layoutlm.embeddings.word_embeddings + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + bbox: 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, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | SequenceClassifierOutput: + r""" + bbox (`torch.LongTensor` of shape `(batch_size, sequence_length, 4)`, *optional*): + Bounding boxes of each input sequence tokens. Selected in the range `[0, + config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1) + format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1, + y1) represents the position of the lower right corner. See [Overview](#Overview) for normalization. + 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 + >>> from transformers import AutoTokenizer, LayoutLMForSequenceClassification + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased") + >>> model = LayoutLMForSequenceClassification.from_pretrained("microsoft/layoutlm-base-uncased") + + >>> words = ["Hello", "world"] + >>> normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782] + + >>> token_boxes = [] + >>> for word, box in zip(words, normalized_word_boxes): + ... word_tokens = tokenizer.tokenize(word) + ... token_boxes.extend([box] * len(word_tokens)) + >>> # add bounding boxes of cls + sep tokens + >>> token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]] + + >>> encoding = tokenizer(" ".join(words), return_tensors="pt") + >>> input_ids = encoding["input_ids"] + >>> attention_mask = encoding["attention_mask"] + >>> token_type_ids = encoding["token_type_ids"] + >>> bbox = torch.tensor([token_boxes]) + >>> sequence_label = torch.tensor([1]) + + >>> outputs = model( + ... input_ids=input_ids, + ... bbox=bbox, + ... attention_mask=attention_mask, + ... token_type_ids=token_type_ids, + ... labels=sequence_label, + ... ) + + >>> loss = outputs.loss + >>> logits = outputs.logits + ```""" + outputs = self.layoutlm( + input_ids=input_ids, + bbox=bbox, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **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( + custom_intro=""" + LayoutLM Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for + sequence labeling (information extraction) tasks such as the [FUNSD](https://guillaumejaume.github.io/FUNSD/) + dataset and the [SROIE](https://rrc.cvc.uab.es/?ch=13) dataset. + """ +) +class LayoutLMForTokenClassification(LayoutLMPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.layoutlm = LayoutLMModel(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() + + def get_input_embeddings(self): + return self.layoutlm.embeddings.word_embeddings + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + bbox: 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, + labels: torch.LongTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | TokenClassifierOutput: + r""" + bbox (`torch.LongTensor` of shape `(batch_size, sequence_length, 4)`, *optional*): + Bounding boxes of each input sequence tokens. Selected in the range `[0, + config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1) + format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1, + y1) represents the position of the lower right corner. See [Overview](#Overview) for normalization. + 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]`. + + Examples: + + ```python + >>> from transformers import AutoTokenizer, LayoutLMForTokenClassification + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("microsoft/layoutlm-base-uncased") + >>> model = LayoutLMForTokenClassification.from_pretrained("microsoft/layoutlm-base-uncased") + + >>> words = ["Hello", "world"] + >>> normalized_word_boxes = [637, 773, 693, 782], [698, 773, 733, 782] + + >>> token_boxes = [] + >>> for word, box in zip(words, normalized_word_boxes): + ... word_tokens = tokenizer.tokenize(word) + ... token_boxes.extend([box] * len(word_tokens)) + >>> # add bounding boxes of cls + sep tokens + >>> token_boxes = [[0, 0, 0, 0]] + token_boxes + [[1000, 1000, 1000, 1000]] + + >>> encoding = tokenizer(" ".join(words), return_tensors="pt") + >>> input_ids = encoding["input_ids"] + >>> attention_mask = encoding["attention_mask"] + >>> token_type_ids = encoding["token_type_ids"] + >>> bbox = torch.tensor([token_boxes]) + >>> token_labels = torch.tensor([1, 1, 0, 0]).unsqueeze(0) # batch size of 1 + + >>> outputs = model( + ... input_ids=input_ids, + ... bbox=bbox, + ... attention_mask=attention_mask, + ... token_type_ids=token_type_ids, + ... labels=token_labels, + ... ) + + >>> loss = outputs.loss + >>> logits = outputs.logits + ```""" + outputs = self.layoutlm( + input_ids=input_ids, + bbox=bbox, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **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 LayoutLMForQuestionAnswering(LayoutLMPreTrainedModel): + def __init__(self, config, has_visual_segment_embedding=True): + r""" + has_visual_segment_embedding (`bool`, *optional*, defaults to `True`): + Whether or not to add visual segment embeddings. + """ + super().__init__(config) + self.num_labels = config.num_labels + + self.layoutlm = LayoutLMModel(config) + self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.layoutlm.embeddings.word_embeddings + + @can_return_tuple + @auto_docstring + def forward( + self, + input_ids: torch.LongTensor | None = None, + bbox: 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: Unpack[TransformersKwargs], + ) -> tuple | QuestionAnsweringModelOutput: + r""" + bbox (`torch.LongTensor` of shape `(batch_size, sequence_length, 4)`, *optional*): + Bounding boxes of each input sequence tokens. Selected in the range `[0, + config.max_2d_position_embeddings-1]`. Each bounding box should be a normalized version in (x0, y0, x1, y1) + format, where (x0, y0) corresponds to the position of the upper left corner in the bounding box, and (x1, + y1) represents the position of the lower right corner. See [Overview](#Overview) for normalization. + + Example: + + In the example below, we prepare a question + context pair for the LayoutLM model. It will give us a prediction + of what it thinks the answer is (the span of the answer within the texts parsed from the image). + + ```python + >>> from transformers import AutoTokenizer, LayoutLMForQuestionAnswering + >>> from datasets import load_dataset + >>> import torch + + >>> tokenizer = AutoTokenizer.from_pretrained("impira/layoutlm-document-qa", add_prefix_space=True) + >>> model = LayoutLMForQuestionAnswering.from_pretrained("impira/layoutlm-document-qa", revision="1e3ebac") + + >>> dataset = load_dataset("nielsr/funsd", split="train") + >>> example = dataset[0] + >>> question = "what's his name?" + >>> words = example["words"] + >>> boxes = example["bboxes"] + + >>> encoding = tokenizer( + ... question.split(), words, is_split_into_words=True, return_token_type_ids=True, return_tensors="pt" + ... ) + >>> bbox = [] + >>> for i, s, w in zip(encoding.input_ids[0], encoding.sequence_ids(0), encoding.word_ids(0)): + ... if s == 1: + ... bbox.append(boxes[w]) + ... elif i == tokenizer.sep_token_id: + ... bbox.append([1000] * 4) + ... else: + ... bbox.append([0] * 4) + >>> encoding["bbox"] = torch.tensor([bbox]) + + >>> word_ids = encoding.word_ids(0) + >>> outputs = model(**encoding) + >>> loss = outputs.loss + >>> start_scores = outputs.start_logits + >>> end_scores = outputs.end_logits + >>> start, end = word_ids[start_scores.argmax(-1)], word_ids[end_scores.argmax(-1)] + >>> print(" ".join(words[start : end + 1])) + M. Hamann P. Harper, P. Martinez + ```""" + + outputs = self.layoutlm( + input_ids=input_ids, + bbox=bbox, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **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__ = [ + "LayoutLMForMaskedLM", + "LayoutLMForSequenceClassification", + "LayoutLMForTokenClassification", + "LayoutLMForQuestionAnswering", + "LayoutLMModel", + "LayoutLMPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/led/__init__.py b/third_party/transformers/src/transformers/models/led/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9609d25e71c9c611e3169a67cc2a3c7a186472c --- /dev/null +++ b/third_party/transformers/src/transformers/models/led/__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 LEDTokenizer + from .configuration_led import * + from .modeling_led 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/led/configuration_led.py b/third_party/transformers/src/transformers/models/led/configuration_led.py new file mode 100644 index 0000000000000000000000000000000000000000..69fc466b2169abbd53791f4a15b4781904793624 --- /dev/null +++ b/third_party/transformers/src/transformers/models/led/configuration_led.py @@ -0,0 +1,86 @@ +# Copyright 2021 Iz Beltagy, Matthew E. Peters, Arman Cohan 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. +"""LED model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="allenai/led-base-16384") +@strict +class LEDConfig(PreTrainedConfig): + r""" + max_encoder_position_embeddings (`int`, *optional*, defaults to 16384): + The maximum sequence length that the encoder might ever be used with. + max_decoder_position_embeddings (`int`, *optional*, defaults to 16384): + The maximum sequence length that the decoder might ever be used with. + attention_window (`int` or `list[int]`, *optional*, defaults to 512): + Size of an attention window around each token. If an `int`, use the same size for all layers. To specify a + different window size for each layer, use a `list[int]` where `len(attention_window) == num_hidden_layers`. + + Example: + + ```python + >>> from transformers import LEDModel, LEDConfig + + >>> # Initializing a LED allenai/led-base-16384 style configuration + >>> configuration = LEDConfig() + + >>> # Initializing a model from the allenai/led-base-16384 style configuration + >>> model = LEDModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "led" + attribute_map = { + "num_attention_heads": "encoder_attention_heads", + "hidden_size": "d_model", + "attention_probs_dropout_prob": "attention_dropout", + "initializer_range": "init_std", + "num_hidden_layers": "encoder_layers", + } + + vocab_size: int = 50265 + max_encoder_position_embeddings: int = 16384 + max_decoder_position_embeddings: int = 1024 + encoder_layers: int = 12 + encoder_ffn_dim: int = 4096 + encoder_attention_heads: int = 16 + decoder_layers: int = 12 + decoder_ffn_dim: int = 4096 + 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 = 1024 + 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 = 2 + classifier_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 + attention_window: list[int] | int = 512 + tie_word_embeddings: bool = True + + +__all__ = ["LEDConfig"] diff --git a/third_party/transformers/src/transformers/models/led/modeling_led.py b/third_party/transformers/src/transformers/models/led/modeling_led.py new file mode 100644 index 0000000000000000000000000000000000000000..6fb1baea2bf57de1178a981139c3c9d61f6bb766 --- /dev/null +++ b/third_party/transformers/src/transformers/models/led/modeling_led.py @@ -0,0 +1,2202 @@ +# Copyright 2021 Iz Beltagy, Matthew E. Peters, Arman Cohan 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 LED model.""" + +import math +from dataclasses import dataclass + +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_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPastAndCrossAttentions +from ...modeling_utils import PreTrainedModel +from ...utils import ModelOutput, auto_docstring, logging +from .configuration_led import LEDConfig + + +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("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 + + +def _prepare_4d_attention_mask_inverted(mask: torch.Tensor, dtype: torch.dtype, tgt_len: int | None = None): + """ + Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`. + """ + bsz, src_len = mask.size() + tgt_len = tgt_len if tgt_len is not None else src_len + + expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype) + + inverted_mask = 1.0 - expanded_mask + expanded_attention_mask = inverted_mask.masked_fill(inverted_mask.bool(), torch.finfo(dtype).min) + + # make sure that global_attn_mask is positive + expanded_attention_mask = expanded_attention_mask * inverted_mask + + return expanded_attention_mask + + +class LEDLearnedPositionalEmbedding(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): + """`input_ids_shape` is expected to be [bsz x seqlen].""" + bsz, seq_len = input_ids_shape[:2] + positions = torch.arange( + past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device + ) + return super().forward(positions) + + +# Copied from transformers.models.longformer.modeling_longformer.LongformerSelfAttention with Longformer->LEDEncoder +class LEDEncoderSelfAttention(nn.Module): + def __init__(self, config, layer_id): + super().__init__() + if config.hidden_size % config.num_attention_heads != 0: + 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_heads = config.num_attention_heads + self.head_dim = int(config.hidden_size / config.num_attention_heads) + self.embed_dim = config.hidden_size + + self.query = nn.Linear(config.hidden_size, self.embed_dim) + self.key = nn.Linear(config.hidden_size, self.embed_dim) + self.value = nn.Linear(config.hidden_size, self.embed_dim) + + # separate projection layers for tokens with global attention + self.query_global = nn.Linear(config.hidden_size, self.embed_dim) + self.key_global = nn.Linear(config.hidden_size, self.embed_dim) + self.value_global = nn.Linear(config.hidden_size, self.embed_dim) + + self.dropout = config.attention_probs_dropout_prob + + self.layer_id = layer_id + attention_window = config.attention_window[self.layer_id] + assert attention_window % 2 == 0, ( + f"`attention_window` for layer {self.layer_id} has to be an even value. Given {attention_window}" + ) + assert attention_window > 0, ( + f"`attention_window` for layer {self.layer_id} has to be positive. Given {attention_window}" + ) + + self.one_sided_attn_window_size = attention_window // 2 + + self.config = config + + def forward( + self, + hidden_states, + attention_mask=None, + is_index_masked=None, + is_index_global_attn=None, + is_global_attn=None, + output_attentions=False, + ): + """ + [`LEDEncoderSelfAttention`] expects *len(hidden_states)* to be multiple of *attention_window*. Padding to + *attention_window* happens in [`LEDEncoderModel.forward`] to avoid redoing the padding on each layer. + + The *attention_mask* is changed in [`LEDEncoderModel.forward`] from 0, 1, 2 to: + + - -10000: no attention + - 0: local attention + - +10000: global attention + """ + hidden_states = hidden_states.transpose(0, 1) + + # project hidden states + query_vectors = self.query(hidden_states) + key_vectors = self.key(hidden_states) + value_vectors = self.value(hidden_states) + + seq_len, batch_size, embed_dim = hidden_states.size() + assert embed_dim == self.embed_dim, ( + f"hidden_states should have embed_dim = {self.embed_dim}, but has {embed_dim}" + ) + + # normalize query + query_vectors /= math.sqrt(self.head_dim) + + query_vectors = query_vectors.view(seq_len, batch_size, self.num_heads, self.head_dim).transpose(0, 1) + key_vectors = key_vectors.view(seq_len, batch_size, self.num_heads, self.head_dim).transpose(0, 1) + + attn_scores = self._sliding_chunks_query_key_matmul( + query_vectors, key_vectors, self.one_sided_attn_window_size + ) + + # values to pad for attention probs + remove_from_windowed_attention_mask = (attention_mask != 0)[:, :, None, None] + + # cast to fp32/fp16 then replace 1's with -inf + float_mask = remove_from_windowed_attention_mask.type_as(query_vectors).masked_fill( + remove_from_windowed_attention_mask, torch.finfo(query_vectors.dtype).min + ) + # diagonal mask with zeros everywhere and -inf inplace of padding + diagonal_mask = self._sliding_chunks_query_key_matmul( + float_mask.new_ones(size=float_mask.size()), float_mask, self.one_sided_attn_window_size + ) + + # pad local attention probs + attn_scores += diagonal_mask + + assert list(attn_scores.size()) == [ + batch_size, + seq_len, + self.num_heads, + self.one_sided_attn_window_size * 2 + 1, + ], ( + f"local_attn_probs should be of size ({batch_size}, {seq_len}, {self.num_heads}," + f" {self.one_sided_attn_window_size * 2 + 1}), but is of size {attn_scores.size()}" + ) + + # compute local attention probs from global attention keys and contact over window dim + if is_global_attn: + # compute global attn indices required through out forward fn + ( + max_num_global_attn_indices, + is_index_global_attn_nonzero, + is_local_index_global_attn_nonzero, + is_local_index_no_global_attn_nonzero, + ) = self._get_global_attn_indices(is_index_global_attn) + # calculate global attn probs from global key + + global_key_attn_scores = self._concat_with_global_key_attn_probs( + query_vectors=query_vectors, + key_vectors=key_vectors, + max_num_global_attn_indices=max_num_global_attn_indices, + is_index_global_attn_nonzero=is_index_global_attn_nonzero, + is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero, + is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero, + ) + # concat to local_attn_probs + # (batch_size, seq_len, num_heads, extra attention count + 2*window+1) + attn_scores = torch.cat((global_key_attn_scores, attn_scores), dim=-1) + + # free memory + del global_key_attn_scores + + attn_probs = nn.functional.softmax( + attn_scores, dim=-1, dtype=torch.float32 + ) # use fp32 for numerical stability + + # softmax sometimes inserts NaN if all positions are masked, replace them with 0 + attn_probs = torch.masked_fill(attn_probs, is_index_masked[:, :, None, None], 0.0) + attn_probs = attn_probs.type_as(attn_scores) + + # free memory + del attn_scores + + # apply dropout + attn_probs = nn.functional.dropout(attn_probs, p=self.dropout, training=self.training) + + value_vectors = value_vectors.view(seq_len, batch_size, self.num_heads, self.head_dim).transpose(0, 1) + + # compute local attention output with global attention value and add + if is_global_attn: + # compute sum of global and local attn + attn_output = self._compute_attn_output_with_global_indices( + value_vectors=value_vectors, + attn_probs=attn_probs, + max_num_global_attn_indices=max_num_global_attn_indices, + is_index_global_attn_nonzero=is_index_global_attn_nonzero, + is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero, + ) + else: + # compute local attn only + attn_output = self._sliding_chunks_matmul_attn_probs_value( + attn_probs, value_vectors, self.one_sided_attn_window_size + ) + + assert attn_output.size() == (batch_size, seq_len, self.num_heads, self.head_dim), "Unexpected size" + attn_output = attn_output.transpose(0, 1).reshape(seq_len, batch_size, embed_dim).contiguous() + + # compute value for global attention and overwrite to attention output + # TODO: remove the redundant computation + if is_global_attn: + global_attn_output, global_attn_probs = self._compute_global_attn_output_from_hidden( + hidden_states=hidden_states, + max_num_global_attn_indices=max_num_global_attn_indices, + is_local_index_global_attn_nonzero=is_local_index_global_attn_nonzero, + is_index_global_attn_nonzero=is_index_global_attn_nonzero, + is_local_index_no_global_attn_nonzero=is_local_index_no_global_attn_nonzero, + is_index_masked=is_index_masked, + ) + + # get only non zero global attn output + nonzero_global_attn_output = global_attn_output[ + is_local_index_global_attn_nonzero[0], :, is_local_index_global_attn_nonzero[1] + ] + + # overwrite values with global attention + attn_output[is_index_global_attn_nonzero[::-1]] = nonzero_global_attn_output.view( + len(is_local_index_global_attn_nonzero[0]), -1 + ) + # The attention weights for tokens with global attention are + # just filler values, they were never used to compute the output. + # Fill with 0 now, the correct values are in 'global_attn_probs'. + attn_probs[is_index_global_attn_nonzero] = 0 + + outputs = (attn_output.transpose(0, 1),) + + if output_attentions: + outputs += (attn_probs,) + + return outputs + (global_attn_probs,) if (is_global_attn and output_attentions) else outputs + + @staticmethod + def _pad_and_transpose_last_two_dims(hidden_states_padded, padding): + """pads rows and then flips rows and columns""" + hidden_states_padded = nn.functional.pad( + hidden_states_padded, padding + ) # padding value is not important because it will be overwritten + hidden_states_padded = hidden_states_padded.view( + *hidden_states_padded.size()[:-2], hidden_states_padded.size(-1), hidden_states_padded.size(-2) + ) + return hidden_states_padded + + @staticmethod + def _pad_and_diagonalize(chunked_hidden_states): + """ + shift every row 1 step right, converting columns into diagonals. + + Example: + + ```python + chunked_hidden_states: [ + 0.4983, + 2.6918, + -0.0071, + 1.0492, + -1.8348, + 0.7672, + 0.2986, + 0.0285, + -0.7584, + 0.4206, + -0.0405, + 0.1599, + 2.0514, + -1.1600, + 0.5372, + 0.2629, + ] + window_overlap = num_rows = 4 + ``` + + (pad & diagonalize) => [ 0.4983, 2.6918, -0.0071, 1.0492, 0.0000, 0.0000, 0.0000 + 0.0000, -1.8348, 0.7672, 0.2986, 0.0285, 0.0000, 0.0000 0.0000, 0.0000, -0.7584, 0.4206, + -0.0405, 0.1599, 0.0000 0.0000, 0.0000, 0.0000, 2.0514, -1.1600, 0.5372, 0.2629 ] + """ + total_num_heads, num_chunks, window_overlap, hidden_dim = chunked_hidden_states.size() + chunked_hidden_states = nn.functional.pad( + chunked_hidden_states, (0, window_overlap + 1) + ) # total_num_heads x num_chunks x window_overlap x (hidden_dim+window_overlap+1). Padding value is not important because it'll be overwritten + chunked_hidden_states = chunked_hidden_states.view( + total_num_heads, num_chunks, -1 + ) # total_num_heads x num_chunks x window_overlap*window_overlap+window_overlap + chunked_hidden_states = chunked_hidden_states[ + :, :, :-window_overlap + ] # total_num_heads x num_chunks x window_overlap*window_overlap + chunked_hidden_states = chunked_hidden_states.view( + total_num_heads, num_chunks, window_overlap, window_overlap + hidden_dim + ) + chunked_hidden_states = chunked_hidden_states[:, :, :, :-1] + return chunked_hidden_states + + @staticmethod + def _chunk(hidden_states, window_overlap, onnx_export: bool = False): + """convert into overlapping chunks. Chunk size = 2w, overlap size = w""" + if not onnx_export: + # non-overlapping chunks of size = 2w + hidden_states = hidden_states.view( + hidden_states.size(0), + torch.div(hidden_states.size(1), (window_overlap * 2), rounding_mode="trunc"), + window_overlap * 2, + hidden_states.size(2), + ) + # use `as_strided` to make the chunks overlap with an overlap size = window_overlap + chunk_size = list(hidden_states.size()) + chunk_size[1] = chunk_size[1] * 2 - 1 + + chunk_stride = list(hidden_states.stride()) + chunk_stride[1] = chunk_stride[1] // 2 + return hidden_states.as_strided(size=chunk_size, stride=chunk_stride) + + # When exporting to ONNX, use this separate logic + # have to use slow implementation since as_strided, unfold and 2d-tensor indexing aren't supported (yet) in ONNX export + + # TODO replace this with + # > return hidden_states.unfold(dimension=1, size=window_overlap * 2, step=window_overlap).transpose(2, 3) + # once `unfold` is supported + # the case hidden_states.size(1) == window_overlap * 2 can also simply return hidden_states.unsqueeze(1), but that's control flow + + chunk_size = [ + hidden_states.size(0), + torch.div(hidden_states.size(1), window_overlap, rounding_mode="trunc") - 1, + window_overlap * 2, + hidden_states.size(2), + ] + + overlapping_chunks = torch.empty(chunk_size, device=hidden_states.device) + for chunk in range(chunk_size[1]): + overlapping_chunks[:, chunk, :, :] = hidden_states[ + :, chunk * window_overlap : chunk * window_overlap + 2 * window_overlap, : + ] + return overlapping_chunks + + @staticmethod + def _mask_invalid_locations(input_tensor, affected_seq_len) -> torch.Tensor: + beginning_mask_2d = input_tensor.new_ones(affected_seq_len, affected_seq_len + 1).tril().flip(dims=[0]) + beginning_mask = beginning_mask_2d[None, :, None, :] + ending_mask = beginning_mask.flip(dims=(1, 3)) + beginning_input = input_tensor[:, :affected_seq_len, :, : affected_seq_len + 1] + beginning_mask = beginning_mask.expand(beginning_input.size()) + input_tensor[:, :affected_seq_len, :, : affected_seq_len + 1] = torch.full_like( + beginning_input, -float("inf") + ).where(beginning_mask.bool(), beginning_input) + ending_input = input_tensor[:, -affected_seq_len:, :, -(affected_seq_len + 1) :] + ending_mask = ending_mask.expand(ending_input.size()) + input_tensor[:, -affected_seq_len:, :, -(affected_seq_len + 1) :] = torch.full_like( + ending_input, -float("inf") + ).where(ending_mask.bool(), ending_input) + + def _sliding_chunks_query_key_matmul(self, query: torch.Tensor, key: torch.Tensor, window_overlap: int): + """ + Matrix multiplication of query and key tensors using with a sliding window attention pattern. This + implementation splits the input into overlapping chunks of size 2w (e.g. 512 for pretrained LEDEncoder) with an + overlap of size window_overlap + """ + batch_size, seq_len, num_heads, head_dim = query.size() + assert seq_len % (window_overlap * 2) == 0, ( + f"Sequence length should be multiple of {window_overlap * 2}. Given {seq_len}" + ) + assert query.size() == key.size() + + chunks_count = torch.div(seq_len, window_overlap, rounding_mode="trunc") - 1 + + # group batch_size and num_heads dimensions into one, then chunk seq_len into chunks of size window_overlap * 2 + query = query.transpose(1, 2).reshape(batch_size * num_heads, seq_len, head_dim) + key = key.transpose(1, 2).reshape(batch_size * num_heads, seq_len, head_dim) + + query = self._chunk(query, window_overlap, getattr(self.config, "onnx_export", False)) + key = self._chunk(key, window_overlap, getattr(self.config, "onnx_export", False)) + + # matrix multiplication + # bcxd: batch_size * num_heads x chunks x 2window_overlap x head_dim + # bcyd: batch_size * num_heads x chunks x 2window_overlap x head_dim + # bcxy: batch_size * num_heads x chunks x 2window_overlap x 2window_overlap + diagonal_chunked_attention_scores = torch.einsum("bcxd,bcyd->bcxy", (query, key)) # multiply + + # convert diagonals into columns + diagonal_chunked_attention_scores = self._pad_and_transpose_last_two_dims( + diagonal_chunked_attention_scores, padding=(0, 0, 0, 1) + ) + + # allocate space for the overall attention matrix where the chunks are combined. The last dimension + # has (window_overlap * 2 + 1) columns. The first (window_overlap) columns are the window_overlap lower triangles (attention from a word to + # window_overlap previous words). The following column is attention score from each word to itself, then + # followed by window_overlap columns for the upper triangle. + + diagonal_attention_scores = diagonal_chunked_attention_scores.new_zeros( + (batch_size * num_heads, chunks_count + 1, window_overlap, window_overlap * 2 + 1) + ) + + # copy parts from diagonal_chunked_attention_scores into the combined matrix of attentions + # - copying the main diagonal and the upper triangle + diagonal_attention_scores[:, :-1, :, window_overlap:] = diagonal_chunked_attention_scores[ + :, :, :window_overlap, : window_overlap + 1 + ] + diagonal_attention_scores[:, -1, :, window_overlap:] = diagonal_chunked_attention_scores[ + :, -1, window_overlap:, : window_overlap + 1 + ] + # - copying the lower triangle + diagonal_attention_scores[:, 1:, :, :window_overlap] = diagonal_chunked_attention_scores[ + :, :, -(window_overlap + 1) : -1, window_overlap + 1 : + ] + + diagonal_attention_scores[:, 0, 1:window_overlap, 1:window_overlap] = diagonal_chunked_attention_scores[ + :, 0, : window_overlap - 1, 1 - window_overlap : + ] + + # separate batch_size and num_heads dimensions again + diagonal_attention_scores = diagonal_attention_scores.view( + batch_size, num_heads, seq_len, 2 * window_overlap + 1 + ).transpose(2, 1) + + self._mask_invalid_locations(diagonal_attention_scores, window_overlap) + return diagonal_attention_scores + + def _sliding_chunks_matmul_attn_probs_value( + self, attn_probs: torch.Tensor, value: torch.Tensor, window_overlap: int + ): + """ + Same as _sliding_chunks_query_key_matmul but for attn_probs and value tensors. Returned tensor will be of the + same shape as `attn_probs` + """ + batch_size, seq_len, num_heads, head_dim = value.size() + + assert seq_len % (window_overlap * 2) == 0 + assert attn_probs.size()[:3] == value.size()[:3] + assert attn_probs.size(3) == 2 * window_overlap + 1 + chunks_count = torch.div(seq_len, window_overlap, rounding_mode="trunc") - 1 + # group batch_size and num_heads dimensions into one, then chunk seq_len into chunks of size 2 window overlap + + chunked_attn_probs = attn_probs.transpose(1, 2).reshape( + batch_size * num_heads, + torch.div(seq_len, window_overlap, rounding_mode="trunc"), + window_overlap, + 2 * window_overlap + 1, + ) + + # group batch_size and num_heads dimensions into one + value = value.transpose(1, 2).reshape(batch_size * num_heads, seq_len, head_dim) + + # pad seq_len with w at the beginning of the sequence and another window overlap at the end + padded_value = nn.functional.pad(value, (0, 0, window_overlap, window_overlap), value=-1) + + # chunk padded_value into chunks of size 3 window overlap and an overlap of size window overlap + chunked_value_size = (batch_size * num_heads, chunks_count + 1, 3 * window_overlap, head_dim) + chunked_value_stride = padded_value.stride() + chunked_value_stride = ( + chunked_value_stride[0], + window_overlap * chunked_value_stride[1], + chunked_value_stride[1], + chunked_value_stride[2], + ) + chunked_value = padded_value.as_strided(size=chunked_value_size, stride=chunked_value_stride) + + chunked_attn_probs = self._pad_and_diagonalize(chunked_attn_probs) + + context = torch.einsum("bcwd,bcdh->bcwh", (chunked_attn_probs, chunked_value)) + return context.view(batch_size, num_heads, seq_len, head_dim).transpose(1, 2) + + @staticmethod + def _get_global_attn_indices(is_index_global_attn): + """compute global attn indices required throughout forward pass""" + # helper variable + num_global_attn_indices = is_index_global_attn.long().sum(dim=1) + + # max number of global attn indices in batch + max_num_global_attn_indices = num_global_attn_indices.max() + + # indices of global attn + is_index_global_attn_nonzero = is_index_global_attn.nonzero(as_tuple=True) + + # helper variable + is_local_index_global_attn = torch.arange( + max_num_global_attn_indices, device=is_index_global_attn.device + ) < num_global_attn_indices.unsqueeze(dim=-1) + + # location of the non-padding values within global attention indices + is_local_index_global_attn_nonzero = is_local_index_global_attn.nonzero(as_tuple=True) + + # location of the padding values within global attention indices + is_local_index_no_global_attn_nonzero = (is_local_index_global_attn == 0).nonzero(as_tuple=True) + return ( + max_num_global_attn_indices, + is_index_global_attn_nonzero, + is_local_index_global_attn_nonzero, + is_local_index_no_global_attn_nonzero, + ) + + def _concat_with_global_key_attn_probs( + self, + key_vectors, + query_vectors, + max_num_global_attn_indices, + is_index_global_attn_nonzero, + is_local_index_global_attn_nonzero, + is_local_index_no_global_attn_nonzero, + ): + batch_size = key_vectors.shape[0] + + # create only global key vectors + key_vectors_only_global = key_vectors.new_zeros( + batch_size, max_num_global_attn_indices, self.num_heads, self.head_dim + ) + + key_vectors_only_global[is_local_index_global_attn_nonzero] = key_vectors[is_index_global_attn_nonzero] + + # (batch_size, seq_len, num_heads, max_num_global_attn_indices) + attn_probs_from_global_key = torch.einsum("blhd,bshd->blhs", (query_vectors, key_vectors_only_global)) + + # need to transpose since ONNX export only supports consecutive indexing: https://pytorch.org/docs/stable/onnx.html#writes-sets + attn_probs_from_global_key = attn_probs_from_global_key.transpose(1, 3) + attn_probs_from_global_key[ + is_local_index_no_global_attn_nonzero[0], is_local_index_no_global_attn_nonzero[1], :, : + ] = torch.finfo(attn_probs_from_global_key.dtype).min + attn_probs_from_global_key = attn_probs_from_global_key.transpose(1, 3) + + return attn_probs_from_global_key + + def _compute_attn_output_with_global_indices( + self, + value_vectors, + attn_probs, + max_num_global_attn_indices, + is_index_global_attn_nonzero, + is_local_index_global_attn_nonzero, + ): + batch_size = attn_probs.shape[0] + + # cut local attn probs to global only + attn_probs_only_global = attn_probs.narrow(-1, 0, max_num_global_attn_indices) + # get value vectors for global only + value_vectors_only_global = value_vectors.new_zeros( + batch_size, max_num_global_attn_indices, self.num_heads, self.head_dim + ) + value_vectors_only_global[is_local_index_global_attn_nonzero] = value_vectors[is_index_global_attn_nonzero] + + # use `matmul` because `einsum` crashes sometimes with fp16 + # attn = torch.einsum('blhs,bshd->blhd', (selected_attn_probs, selected_v)) + # compute attn output only global + attn_output_only_global = torch.matmul( + attn_probs_only_global.transpose(1, 2).clone(), value_vectors_only_global.transpose(1, 2).clone() + ).transpose(1, 2) + + # reshape attn probs + attn_probs_without_global = attn_probs.narrow( + -1, max_num_global_attn_indices, attn_probs.size(-1) - max_num_global_attn_indices + ).contiguous() + + # compute attn output with global + attn_output_without_global = self._sliding_chunks_matmul_attn_probs_value( + attn_probs_without_global, value_vectors, self.one_sided_attn_window_size + ) + return attn_output_only_global + attn_output_without_global + + def _compute_global_attn_output_from_hidden( + self, + hidden_states, + max_num_global_attn_indices, + is_local_index_global_attn_nonzero, + is_index_global_attn_nonzero, + is_local_index_no_global_attn_nonzero, + is_index_masked, + ): + seq_len, batch_size = hidden_states.shape[:2] + + # prepare global hidden states + global_attn_hidden_states = hidden_states.new_zeros(max_num_global_attn_indices, batch_size, self.embed_dim) + global_attn_hidden_states[is_local_index_global_attn_nonzero[::-1]] = hidden_states[ + is_index_global_attn_nonzero[::-1] + ] + + # global key, query, value + global_query_vectors_only_global = self.query_global(global_attn_hidden_states) + global_key_vectors = self.key_global(hidden_states) + global_value_vectors = self.value_global(hidden_states) + + # normalize + global_query_vectors_only_global /= math.sqrt(self.head_dim) + + # reshape + global_query_vectors_only_global = ( + global_query_vectors_only_global.contiguous() + .view(max_num_global_attn_indices, batch_size * self.num_heads, self.head_dim) + .transpose(0, 1) + ) # (batch_size * self.num_heads, max_num_global_attn_indices, head_dim) + global_key_vectors = ( + global_key_vectors.contiguous().view(-1, batch_size * self.num_heads, self.head_dim).transpose(0, 1) + ) # batch_size * self.num_heads, seq_len, head_dim) + global_value_vectors = ( + global_value_vectors.contiguous().view(-1, batch_size * self.num_heads, self.head_dim).transpose(0, 1) + ) # batch_size * self.num_heads, seq_len, head_dim) + + # compute attn scores + global_attn_scores = torch.bmm(global_query_vectors_only_global, global_key_vectors.transpose(1, 2)) + + assert list(global_attn_scores.size()) == [ + batch_size * self.num_heads, + max_num_global_attn_indices, + seq_len, + ], ( + "global_attn_scores have the wrong size. Size should be" + f" {(batch_size * self.num_heads, max_num_global_attn_indices, seq_len)}, but is" + f" {global_attn_scores.size()}." + ) + + global_attn_scores = global_attn_scores.view(batch_size, self.num_heads, max_num_global_attn_indices, seq_len) + + # need to transpose since ONNX export only supports consecutive indexing: https://pytorch.org/docs/stable/onnx.html#writes-sets + global_attn_scores = global_attn_scores.transpose(1, 2) + global_attn_scores[ + is_local_index_no_global_attn_nonzero[0], is_local_index_no_global_attn_nonzero[1], :, : + ] = torch.finfo(global_attn_scores.dtype).min + global_attn_scores = global_attn_scores.transpose(1, 2) + + global_attn_scores = global_attn_scores.masked_fill( + is_index_masked[:, None, None, :], + torch.finfo(global_attn_scores.dtype).min, + ) + + global_attn_scores = global_attn_scores.view(batch_size * self.num_heads, max_num_global_attn_indices, seq_len) + + # compute global attn probs + global_attn_probs_float = nn.functional.softmax( + global_attn_scores, dim=-1, dtype=torch.float32 + ) # use fp32 for numerical stability + + global_attn_probs = nn.functional.dropout( + global_attn_probs_float.type_as(global_attn_scores), p=self.dropout, training=self.training + ) + + # global attn output + global_attn_output = torch.bmm(global_attn_probs, global_value_vectors) + + assert list(global_attn_output.size()) == [ + batch_size * self.num_heads, + max_num_global_attn_indices, + self.head_dim, + ], ( + "global_attn_output tensor has the wrong size. Size should be" + f" {(batch_size * self.num_heads, max_num_global_attn_indices, self.head_dim)}, but is" + f" {global_attn_output.size()}." + ) + + global_attn_probs = global_attn_probs.view(batch_size, self.num_heads, max_num_global_attn_indices, seq_len) + global_attn_output = global_attn_output.view( + batch_size, self.num_heads, max_num_global_attn_indices, self.head_dim + ) + return global_attn_output, global_attn_probs + + +class LEDEncoderAttention(nn.Module): + def __init__(self, config, layer_id): + super().__init__() + self.longformer_self_attn = LEDEncoderSelfAttention(config, layer_id=layer_id) + self.output = nn.Linear(config.d_model, config.d_model) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + is_index_masked: torch.Tensor | None = None, + is_index_global_attn: torch.Tensor | None = None, + is_global_attn: bool | None = None, + output_attentions: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]: + """Input shape: Batch x Time x Channel""" + + self_outputs = self.longformer_self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + is_index_masked=is_index_masked, + is_index_global_attn=is_index_global_attn, + is_global_attn=is_global_attn, + output_attentions=output_attentions, + ) + + attn_output = self.output(self_outputs[0]) + outputs = (attn_output,) + self_outputs[1:] + + return outputs + + +class LEDDecoderAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__( + self, + embed_dim: int, + num_heads: int, + dropout: float | None = 0.0, + is_decoder: bool | None = False, + bias: bool | None = True, + layer_idx: bool | 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} and `num_heads`:" + f" {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) + + 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, + output_attentions: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, Cache | 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, embed_dim = hidden_states.size() + + # get query proj + query_states = self.q_proj(hidden_states) * self.scaling + + 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(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: + # save all key/value_states to cache to be re-used for fast auto-regressive generation + 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) + + 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) + + if output_attentions: + # 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, src_len) + attn_weights = attn_weights_reshaped.view(bsz * self.num_heads, tgt_len, src_len) + else: + attn_weights_reshaped = None + + 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) + .transpose(1, 2) + .reshape(bsz, tgt_len, embed_dim) + ) + + attn_output = self.out_proj(attn_output) + + return attn_output, attn_weights_reshaped, past_key_values + + +class LEDEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: LEDConfig, layer_id: int): + super().__init__() + self.embed_dim = config.d_model + self.self_attn = LEDEncoderAttention(config, layer_id) + 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, + is_index_masked=None, + is_index_global_attn=None, + is_global_attn=None, + output_attentions=False, + ): + """ + 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 + attn_outputs = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + is_index_masked=is_index_masked, + is_index_global_attn=is_index_global_attn, + is_global_attn=is_global_attn, + output_attentions=output_attentions, + ) + hidden_states = attn_outputs[0] + 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,) + attn_outputs[1:] + + +class LEDDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: LEDConfig, layer_idx=None): + super().__init__() + self.embed_dim = config.d_model + + self.self_attn = LEDDecoderAttention( + embed_dim=self.embed_dim, + num_heads=config.decoder_attention_heads, + dropout=config.attention_dropout, + is_decoder=True, + 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 = LEDDecoderAttention( + self.embed_dim, + config.decoder_attention_heads, + dropout=config.attention_dropout, + is_decoder=True, + 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, + output_attentions: bool | None = False, + use_cache: bool | None = True, + **kwargs, + ): + """ + 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 + output_attentions (`bool`): Whether the base model outputs attentions. + This requires the attentions tensor to be reshaped in this function. + """ + residual = hidden_states + + # Self-Attention + hidden_states, self_attn_weights, present_key_value = self.self_attn( + hidden_states=hidden_states, + past_key_values=past_key_values, + attention_mask=attention_mask, + output_attentions=output_attentions, + ) + 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 + cross_attn_present_key_value = None + cross_attn_weights = None + if encoder_hidden_states is not None: + residual = hidden_states + + hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn( + hidden_states=hidden_states, + key_value_states=encoder_hidden_states, + attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + output_attentions=output_attentions, + ) + 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) + + outputs = (hidden_states,) + + if output_attentions: + outputs += (self_attn_weights, cross_attn_weights) + + if use_cache: + outputs += (past_key_values,) + + return outputs + + +class LEDClassificationHead(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): + 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 LEDPreTrainedModel(PreTrainedModel): + config: LEDConfig + base_model_prefix = "led" + supports_gradient_checkpointing = True + + @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 + + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, LEDForConditionalGeneration): + init.zeros_(module.final_logits_bias) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for LEDEncoder's outputs, with potential hidden states, local and global attentions. + """ +) +# Copied from transformers.models.longformer.modeling_longformer.LongformerBaseModelOutput with Longformer->LEDEncoder +class LEDEncoderBaseModelOutput(ModelOutput): + r""" + 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, x + + attention_window + 1)`, where `x` is the number of tokens with global attention mask. + + Local attentions weights after the attention softmax, used to compute the weighted average in the + self-attention heads. Those are the attention weights from every token in the sequence to every token with + global attention (first `x` values) and to every token in the attention window (remaining `attention_window + + 1` values). Note that the first `x` values refer to tokens with fixed positions in the text, but the + remaining `attention_window + 1` values refer to tokens with relative positions: the attention weight of a + token to itself is located at index `x + attention_window / 2` and the `attention_window / 2` preceding + (succeeding) values are the attention weights to the `attention_window / 2` preceding (succeeding) tokens. + If the attention window contains a token with global attention, the attention weight at the corresponding + index is set to 0; the value should be accessed from the first `x` attention weights. If a token has global + attention, the attention weights to all other tokens in `attentions` is set to 0, the values should be + accessed from `global_attentions`. + global_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, x)`, + where `x` is the number of tokens with global attention mask. + + Global attentions weights after the attention softmax, used to compute the weighted average in the + self-attention heads. Those are the attention weights from every token with global attention to every token + in the sequence. + """ + + last_hidden_state: torch.FloatTensor + hidden_states: tuple[torch.FloatTensor, ...] | None = None + attentions: tuple[torch.FloatTensor, ...] | None = None + global_attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for model encoder's outputs that also contains : pre-computed hidden states that can speed up sequential + decoding. + """ +) +class LEDSeq2SeqModelOutput(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. + 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 attention blocks) of the decoder that can be + used (see `past_key_values` input) to speed up sequential decoding. + encoder_global_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, x)`, + where `x` is the number of tokens with global attention mask. + + Global attentions weights after the attention softmax, used to compute the weighted average in the + self-attention heads. Those are the attention weights from every token with global attention to every token + in the sequence. + """ + + last_hidden_state: 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 + encoder_global_attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for sequence-to-sequence language models outputs. + """ +) +class LEDSeq2SeqLMOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Language modeling 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). + 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 attention blocks) of the decoder that can be + used (see `past_key_values` input) to speed up sequential decoding. + encoder_global_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, x)`, + where `x` is the number of tokens with global attention mask. + + Global attentions weights after the attention softmax, used to compute the weighted average in the + self-attention heads. Those are the attention weights from every token with global attention to every token + in the sequence. + """ + + loss: torch.FloatTensor | None = None + logits: 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 + encoder_global_attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of sequence-to-sequence sentence classification models. + """ +) +class LEDSeq2SeqSequenceClassifierOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `label` 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). + 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 attention blocks) of the decoder that can be + used (see `past_key_values` input) to speed up sequential decoding. + encoder_global_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, x)`, + where `x` is the number of tokens with global attention mask. + + Global attentions weights after the attention softmax, used to compute the weighted average in the + self-attention heads. Those are the attention weights from every token with global attention to every token + in the sequence. + """ + + loss: torch.FloatTensor | None = None + logits: 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 + encoder_global_attentions: tuple[torch.FloatTensor, ...] | None = None + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for outputs of sequence-to-sequence question answering models. + """ +) +class LEDSeq2SeqQuestionAnsweringModelOutput(ModelOutput): + r""" + loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): + Total span extraction loss is the sum of a Cross-Entropy for the start and end positions. + 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 attention blocks) of the decoder that can be + used (see `past_key_values` input) to speed up sequential decoding. + encoder_global_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, x)`, + where `x` is the number of tokens with global attention mask. + + Global attentions weights after the attention softmax, used to compute the weighted average in the + self-attention heads. Those are the attention weights from every token with global attention to every token + in the sequence. + """ + + loss: torch.FloatTensor | None = None + start_logits: torch.FloatTensor | None = None + end_logits: 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 + encoder_global_attentions: tuple[torch.FloatTensor, ...] | None = None + + +class LEDEncoder(LEDPreTrainedModel): + """ + Transformer encoder consisting of *config.encoder_layers* self-attention layers. Each layer is a + [`LEDEncoderLayer`]. + + Args: + config: LEDConfig + embed_tokens (nn.Embedding): output embedding + """ + + def __init__(self, config: LEDConfig): + 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_encoder_position_embeddings + + if isinstance(config.attention_window, int): + if config.attention_window % 2 != 0: + raise ValueError("`config.attention_window` has to be an even value") + if config.attention_window <= 0: + raise ValueError("`config.attention_window` has to be positive") + config.attention_window = [config.attention_window] * config.num_hidden_layers # one value per layer + else: + if len(config.attention_window) != config.num_hidden_layers: + raise ValueError( + "`len(config.attention_window)` should equal `config.num_hidden_layers`. " + f"Expected {config.num_hidden_layers}, given {len(config.attention_window)}" + ) + + self.embed_tokens = nn.Embedding(config.vocab_size, embed_dim, self.padding_idx) + + self.embed_positions = LEDLearnedPositionalEmbedding( + self.max_source_positions, + embed_dim, + ) + self.layers = nn.ModuleList([LEDEncoderLayer(config, 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() + + def _merge_to_attention_mask(self, attention_mask: torch.Tensor, global_attention_mask: torch.Tensor): + # longformer self-attention expects attention mask to have 0 (no attn), 1 (local attn), 2 (global attn) + # (global_attention_mask + 1) => 1 for local attention, 2 for global attention + # => final attention_mask => 0 for no attention, 1 for local attention 2 for global attention + if attention_mask is not None: + attention_mask = attention_mask * (global_attention_mask + 1) + else: + # simply use `global_attention_mask` as `attention_mask` + # if no `attention_mask` is given + attention_mask = global_attention_mask + 1 + return attention_mask + + def _pad_to_window_size( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + inputs_embeds: torch.Tensor, + pad_token_id: int, + ): + """A helper function to pad tokens and mask to work with implementation of Longformer self-attention.""" + # padding + attention_window = ( + self.config.attention_window + if isinstance(self.config.attention_window, int) + else max(self.config.attention_window) + ) + + if attention_window % 2 != 0: + raise ValueError(f"`attention_window` should be an even value. Given {attention_window}") + input_shape = input_ids.shape if input_ids is not None else inputs_embeds.shape + batch_size, seq_len = input_shape[:2] + + padding_len = (attention_window - seq_len % attention_window) % attention_window + if padding_len > 0: + logger.warning_once( + f"Input ids are automatically padded from {seq_len} to {seq_len + padding_len} to be a multiple of " + f"`config.attention_window`: {attention_window}" + ) + if input_ids is not None: + input_ids = nn.functional.pad(input_ids, (0, padding_len), value=pad_token_id) + if inputs_embeds is not None: + input_ids_padding = inputs_embeds.new_full( + (batch_size, padding_len), + self.config.pad_token_id, + dtype=torch.long, + ) + inputs_embeds_padding = self.embed_tokens(input_ids_padding) + inputs_embeds = torch.cat([inputs_embeds, inputs_embeds_padding], dim=-2) + + attention_mask = nn.functional.pad( + attention_mask, (0, padding_len), value=False + ) # no attention on the padding tokens + + return padding_len, input_ids, attention_mask, inputs_embeds + + def forward( + self, + input_ids=None, + attention_mask=None, + global_attention_mask=None, + inputs_embeds=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + **kwargs, + ): + r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you + provide it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + global_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to decide the attention given on each token, local attention or global attention for the encoder. + Tokens with global attention attends to all other tokens, and all other tokens attend to them. This is + important for task-specific finetuning because it makes the model more flexible at representing the + task. For example, for classification, the token should be given global attention. For QA, all + question tokens should also have global attention. Please refer to the [Longformer + paper](https://huggingface.co/papers/2004.05150) for more details. Mask values selected in `[0, 1]`: + + - 0 for local attention (a sliding window attention), + - 1 for global attention (tokens that attend to all other tokens, and all other tokens attend to them). + 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. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + """ + 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 + + # check input_ids and inputs_embeds + 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 None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + # create default attention_mask + if attention_mask is None: + attention_mask = torch.ones(inputs_embeds.size()[:-1], device=inputs_embeds.device, dtype=torch.long) + + # merge `global_attention_mask` and `attention_mask` + if global_attention_mask is not None: + attention_mask = self._merge_to_attention_mask(attention_mask, global_attention_mask) + + # pad input if necessary + padding_len, input_ids, attention_mask, inputs_embeds = self._pad_to_window_size( + input_ids=input_ids, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + pad_token_id=self.config.pad_token_id, + ) + + # retrieve input_shape + if input_ids is not None: + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + + # convert attention_mask to float + if attention_mask is not None: + # [bsz, seq_len] -> [bsz, seq_len]; 1 -> 0.0; 0 -> "-inf" + attention_mask = _prepare_4d_attention_mask_inverted(attention_mask, inputs_embeds.dtype)[:, 0, 0, :] + + # get masking tensors + is_index_masked = attention_mask < 0 + is_index_global_attn = attention_mask > 0 + is_global_attn = is_index_global_attn.flatten().any().item() + + 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) + + encoder_states = () if output_hidden_states else None + all_attentions = () if output_attentions else None + all_global_attentions = () if (output_attentions and is_global_attn) else None + + for idx, encoder_layer in enumerate(self.layers): + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + dropout_probability = torch.rand([]) + + if self.training and (dropout_probability < self.layerdrop): # skip the layer + layer_outputs = (None, None, None) + else: + layer_outputs = encoder_layer( + hidden_states, + attention_mask=attention_mask, + is_index_masked=is_index_masked, + is_index_global_attn=is_index_global_attn, + is_global_attn=is_global_attn, + output_attentions=output_attentions, + ) + hidden_states = layer_outputs[0] + + if output_attentions: + # bzs x seq_len x num_attn_heads x (num_global_attn + attention_window_len + 1) => bzs x num_attn_heads x seq_len x (num_global_attn + attention_window_len + 1) + all_attentions = all_attentions + (layer_outputs[1].transpose(1, 2),) + + if is_global_attn: + # bzs x num_attn_heads x num_global_attn x seq_len => bzs x num_attn_heads x seq_len x num_global_attn + all_global_attentions = all_global_attentions + (layer_outputs[2].transpose(2, 3),) + + if output_hidden_states: + encoder_states = encoder_states + (hidden_states,) + + # undo padding + if padding_len > 0: + # unpad `hidden_states` because the calling function is expecting a length == input_ids.size(1) + hidden_states = hidden_states[:, :-padding_len] + if output_hidden_states: + encoder_states = tuple(state[:, :-padding_len] for state in encoder_states) + + if output_attentions: + all_attentions = tuple(state[:, :, :-padding_len, :] for state in all_attentions) + + if not return_dict: + return tuple( + v for v in [hidden_states, encoder_states, all_attentions, all_global_attentions] if v is not None + ) + return LEDEncoderBaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=encoder_states, + attentions=all_attentions, + global_attentions=all_global_attentions, + ) + + +class LEDDecoder(LEDPreTrainedModel): + """ + Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`LEDDecoderLayer`] + + Args: + config: LEDConfig + embed_tokens (nn.Embedding): output embedding + """ + + def __init__(self, config: LEDConfig): + 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_decoder_position_embeddings + + self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model, self.padding_idx) + + self.embed_positions = LEDLearnedPositionalEmbedding( + self.max_target_positions, + config.d_model, + ) + self.layers = nn.ModuleList([LEDDecoderLayer(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() + + def forward( + self, + input_ids=None, + attention_mask=None, + global_attention_mask=None, + encoder_hidden_states=None, + encoder_attention_mask=None, + past_key_values=None, + inputs_embeds=None, + use_cache=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + **kwargs, + ): + r""" + Args: + input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): + Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you + provide it. + + Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: + + - 1 for tokens that are **not masked**, + - 0 for tokens that are **masked**. + + [What are attention masks?](../glossary#attention-mask) + global_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to decide the attention given on each token, local attention or global attention. Tokens with + global attention attends to all other tokens, and all other tokens attend to them. This is important + for task-specific finetuning because it makes the model more flexible at representing the task. For + example, for classification, the token should be given global attention. For QA, all question + tokens should also have global attention. Please refer to the [Longformer + paper](https://huggingface.co/papers/2004.05150) for more details. Mask values selected in `[0, 1]`: + + - 0 for local attention (a sliding window attention), + - 1 for global attention (tokens that attend to all other tokens, and all other tokens attend to them). + 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. + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under + returned tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors + for more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. + """ + 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 + + # retrieve input_ids and inputs_embeds + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time") + elif input_ids is not None: + input_shape = input_ids.size() + input_ids = input_ids.view(-1, input_shape[-1]) + elif inputs_embeds is not None: + input_shape = inputs_embeds.size()[:-1] + else: + raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + 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 = 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 + + combined_attention_mask = None + if input_shape[-1] > 1: # only create a causal mask when we go over a single token + combined_attention_mask = create_causal_mask( + config=self.config, + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + past_key_values=past_key_values, + ) + + 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_shape, past_key_values_length) + + 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) + + # decoder layers + all_hidden_states = () if output_hidden_states else None + all_self_attns = () if output_attentions else None + all_cross_attentions = () if output_attentions else None + + for idx, decoder_layer in enumerate(self.layers): + # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description) + if output_hidden_states: + all_hidden_states += (hidden_states,) + if self.training: + dropout_probability = torch.rand([]) + if dropout_probability < self.layerdrop: + continue + + layer_outputs = decoder_layer( + hidden_states, + combined_attention_mask, + encoder_hidden_states, # as a positional argument for gradient checkpointing + encoder_attention_mask=encoder_attention_mask, + past_key_values=past_key_values, + output_attentions=output_attentions, + use_cache=use_cache, + ) + + hidden_states = layer_outputs[0] + if output_attentions: + all_self_attns += (layer_outputs[1],) + all_cross_attentions += (layer_outputs[2],) + + # add hidden states from the last decoder layer + if output_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_attns, all_cross_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_attns, + cross_attentions=all_cross_attentions, + ) + + +@auto_docstring +class LEDModel(LEDPreTrainedModel): + _tied_weights_keys = { + "encoder.embed_tokens.weight": "shared.weight", + "decoder.embed_tokens.weight": "shared.weight", + } + + def __init__(self, config: LEDConfig): + 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 = LEDEncoder(config) + self.decoder = LEDDecoder(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 + + @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: tuple[tuple[torch.FloatTensor]] | None = None, + global_attention_mask: 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, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple[torch.Tensor] | LEDSeq2SeqModelOutput: + 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 [`LedTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + + LED 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`). + 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_led._prepare_decoder_inputs`] and modify + to your needs. See diagram 1 in [the paper](https://huggingface.co/papers/1910.13461) for more information on the + default strategy. + global_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to decide the attention given on each token, local attention or global attention for the encoder. + Tokens with global attention attends to all other tokens, and all other tokens attend to them. This is + important for task-specific finetuning because it makes the model more flexible at representing the task. + For example, for classification, the token should be given global attention. For QA, all question + tokens should also have global attention. Please refer to the [Longformer + paper](https://huggingface.co/papers/2004.05150) for more details. Mask values selected in `[0, 1]`: + + - 0 for local attention (a sliding window attention), + - 1 for global attention (tokens that attend to all other tokens, and all other tokens attend to them). + """ + 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 + + # Using this like Bart, as LED is derived from it. So far + # No checkpoint on the hub exists that uses that in practice. + # https://github.com/huggingface/transformers/blob/ac3cb660cad283163f7c73cad511124e845ca388/src/transformers/models/bart/modeling_bart.py#L1153 + if decoder_input_ids is None and decoder_inputs_embeds is None: + 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 = self.encoder( + input_ids=input_ids, + attention_mask=attention_mask, + global_attention_mask=global_attention_mask, + inputs_embeds=inputs_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + # If the user passed a tuple for encoder_outputs, we wrap it in a LEDEncoderBaseModelOutput when return_dict=False + elif return_dict and not isinstance(encoder_outputs, LEDEncoderBaseModelOutput): + encoder_outputs = LEDEncoderBaseModelOutput( + 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, + global_attentions=encoder_outputs[3] if len(encoder_outputs) > 3 else None, + ) + + # decoder outputs consists of (dec_features, past_key_values, dec_hidden, dec_attn) + decoder_outputs = 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, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + if not return_dict: + return decoder_outputs + encoder_outputs + + return LEDSeq2SeqModelOutput( + 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, + encoder_global_attentions=encoder_outputs.global_attentions, + ) + + +@auto_docstring( + custom_intro=""" + The LED Model with a language modeling head. Can be used for summarization. + """ +) +class LEDForConditionalGeneration(LEDPreTrainedModel, GenerationMixin): + base_model_prefix = "led" + _keys_to_ignore_on_load_missing = ["final_logits_bias"] + _tied_weights_keys = { + "lm_head.weight": "led.shared.weight", + } + + def __init__(self, config: LEDConfig): + super().__init__(config) + self.led = LEDModel(config) + self.register_buffer("final_logits_bias", torch.zeros((1, self.led.shared.num_embeddings))) + self.lm_head = nn.Linear(config.d_model, self.led.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) + + @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: tuple[tuple[torch.FloatTensor]] | None = None, + global_attention_mask: 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, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple[torch.Tensor] | LEDSeq2SeqLMOutput: + 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 [`LedTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + + LED 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`). + 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_led._prepare_decoder_inputs`] and modify + to your needs. See diagram 1 in [the paper](https://huggingface.co/papers/1910.13461) for more information on the + default strategy. + global_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to decide the attention given on each token, local attention or global attention for the encoder. + Tokens with global attention attends to all other tokens, and all other tokens attend to them. This is + important for task-specific finetuning because it makes the model more flexible at representing the task. + For example, for classification, the token should be given global attention. For QA, all question + tokens should also have global attention. Please refer to the [Longformer + paper](https://huggingface.co/papers/2004.05150) for more details. Mask values selected in `[0, 1]`: + + - 0 for local attention (a sliding window attention), + - 1 for global attention (tokens that attend to all other tokens, and all other tokens attend to them). + 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 + >>> import torch + >>> from transformers import AutoTokenizer, LEDForConditionalGeneration + + >>> model = LEDForConditionalGeneration.from_pretrained("allenai/led-large-16384-arxiv") + >>> tokenizer = AutoTokenizer.from_pretrained("allenai/led-large-16384-arxiv") + + >>> ARTICLE_TO_SUMMARIZE = '''Transformers (Vaswani et al., 2017) have achieved state-of-the-art + ... results in a wide range of natural language tasks including generative language modeling + ... (Dai et al., 2019; Radford et al., 2019) and discriminative ... language understanding (Devlin et al., 2019). + ... This success is partly due to the self-attention component which enables the network to capture contextual + ... information from the entire sequence. While powerful, the memory and computational requirements of + ... self-attention grow quadratically with sequence length, making it infeasible (or very expensive) to + ... process long sequences. To address this limitation, we present Longformer, a modified Transformer + ... architecture with a self-attention operation that scales linearly with the sequence length, making it + ... versatile for processing long documents (Fig 1). This is an advantage for natural language tasks such as + ... long document classification, question answering (QA), and coreference resolution, where existing approaches + ... partition or shorten the long context into smaller sequences that fall within the typical 512 token limit + ... of BERT-style pretrained models. Such partitioning could potentially result in loss of important + ... cross-partition information, and to mitigate this problem, existing methods often rely on complex + ... architectures to address such interactions. On the other hand, our proposed Longformer is able to build + ... contextual representations of the entire context using multiple layers of attention, reducing the need for + ... task-specific architectures.''' + >>> inputs = tokenizer.encode(ARTICLE_TO_SUMMARIZE, return_tensors="pt") + + >>> # Global attention on the first token (cf. Beltagy et al. 2020) + >>> global_attention_mask = torch.zeros_like(inputs) + >>> global_attention_mask[:, 0] = 1 + + >>> # Generate Summary + >>> summary_ids = model.generate(inputs, global_attention_mask=global_attention_mask, num_beams=3, max_length=32) + >>> print(tokenizer.decode(summary_ids[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)) + ``` + + Example Conditional generation : + + ```python + >>> from transformers import AutoTokenizer, LEDForConditionalGeneration + + >>> tokenizer = AutoTokenizer.from_pretrained("allenai/led-base-16384") + >>> TXT = "My friends are but they eat too many carbs." + + >>> model = LEDForConditionalGeneration.from_pretrained("allenai/led-base-16384") + >>> input_ids = tokenizer([TXT], return_tensors="pt")["input_ids"] + + >>> prediction = model.generate(input_ids)[0] + >>> print(tokenizer.decode(prediction, skip_special_tokens=True)) + ``` + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + + 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 = self.led( + input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + encoder_outputs=encoder_outputs, + global_attention_mask=global_attention_mask, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + lm_logits = self.lm_head(outputs[0]) + self.final_logits_bias + + masked_lm_loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1)) + + if not return_dict: + output = (lm_logits,) + outputs[1:] + return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output + + return LEDSeq2SeqLMOutput( + 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, + encoder_global_attentions=outputs.encoder_global_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 +class LEDForQuestionAnswering(LEDPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + config.num_labels = 2 + self.num_labels = config.num_labels + + self.led = LEDModel(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.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: tuple[tuple[torch.FloatTensor]] | None = None, + global_attention_mask: 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, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple[torch.Tensor] | LEDSeq2SeqQuestionAnsweringModelOutput: + 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 [`LedTokenizer`]. See [`PreTrainedTokenizer.encode`] and + [`PreTrainedTokenizer.__call__`] for details. + + [What are input IDs?](../glossary#input-ids) + + LED 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`). + 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_led._prepare_decoder_inputs`] and modify + to your needs. See diagram 1 in [the paper](https://huggingface.co/papers/1910.13461) for more information on the + default strategy. + global_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*): + Mask to decide the attention given on each token, local attention or global attention for the encoder. + Tokens with global attention attends to all other tokens, and all other tokens attend to them. This is + important for task-specific finetuning because it makes the model more flexible at representing the task. + For example, for classification, the token should be given global attention. For QA, all question + tokens should also have global attention. Please refer to the [Longformer + paper](https://huggingface.co/papers/2004.05150) for more details. Mask values selected in `[0, 1]`: + + - 0 for local attention (a sliding window attention), + - 1 for global attention (tokens that attend to all other tokens, and all other tokens attend to them). + """ + return_dict = return_dict if return_dict is not None else self.config.return_dict + if start_positions is not None and end_positions is not None: + use_cache = False + + outputs = self.led( + input_ids, + attention_mask=attention_mask, + decoder_input_ids=decoder_input_ids, + decoder_attention_mask=decoder_attention_mask, + global_attention_mask=global_attention_mask, + encoder_outputs=encoder_outputs, + inputs_embeds=inputs_embeds, + decoder_inputs_embeds=decoder_inputs_embeds, + use_cache=use_cache, + 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[1:] + return ((total_loss,) + output) if total_loss is not None else output + + return LEDSeq2SeqQuestionAnsweringModelOutput( + 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, + encoder_global_attentions=outputs.encoder_global_attentions, + ) + + +__all__ = [ + "LEDForConditionalGeneration", + "LEDForQuestionAnswering", + "LEDModel", + "LEDPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/lfm2_moe/__init__.py b/third_party/transformers/src/transformers/models/lfm2_moe/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f2f9c74765bd3eeafdc84a513995c16cceeca025 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2_moe/__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_lfm2_moe import * + from .modeling_lfm2_moe 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/lfm2_moe/configuration_lfm2_moe.py b/third_party/transformers/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..b5ae5c76bddec4b4ffd155eddde2c4fb429a9e5c --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2_moe/configuration_lfm2_moe.py @@ -0,0 +1,84 @@ +# 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 + + +@auto_docstring(checkpoint="LiquidAI/LFM2-8B-A1B") +@strict +class Lfm2MoeConfig(PreTrainedConfig): + r""" + conv_bias (`bool`, *optional*, defaults to `False`): + Whether to use bias in the conv layers. + conv_L_cache (`int`, *optional*, defaults to 3): + L_cache dim in the conv layers. + num_dense_layers (`int`, *optional*, defaults to 2): + Number of dense Lfm2MoeMLP layers in shallow layers(embed->dense->dense->...->dense->moe->moe...->lm_head). + use_expert_bias (`bool`, *optional*, defaults to `True`): + Whether to use the expert bias on the routing weights. + + ```python + >>> from transformers import Lfm2MoeModel, Lfm2MoeConfig + + >>> # Initializing a LFM2 Moe model + >>> configuration = Lfm2MoeConfig() + + >>> # Initializing a model from the LFM2-8B-A1B style configuration + >>> model = Lfm2MoeModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "lfm2_moe" + keys_to_ignore_at_inference = ["past_key_values"] + default_theta = 1000000.0 + + vocab_size: int = 65536 + hidden_size: int = 2048 + intermediate_size: int = 7168 + moe_intermediate_size: int = 1792 + num_hidden_layers: int = 32 + pad_token_id: int | None = 0 + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + tie_word_embeddings: bool = True + rope_parameters: dict | None = None + max_position_embeddings: int = 128_000 + initializer_range: float = 0.02 + use_cache: bool = True + norm_eps: float = 0.00001 + num_attention_heads: int = 32 + num_key_value_heads: int = 8 + conv_bias: bool = False + conv_L_cache: int = 3 + num_dense_layers: int = 2 + num_experts_per_tok: int = 4 + num_experts: int = 32 + use_expert_bias: bool = True + routed_scaling_factor: float = 1.0 + norm_topk_prob: bool = True + layer_types: list[str] | None = None + + def __post_init__(self, **kwargs): + self.tie_word_embeddings = kwargs.pop("tie_embedding", self.tie_word_embeddings) + super().__post_init__(**kwargs) + + +__all__ = ["Lfm2MoeConfig"] diff --git a/third_party/transformers/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py b/third_party/transformers/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..0369ae31b8ae8990d965288e961fc0c61870e455 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2_moe/modeling_lfm2_moe.py @@ -0,0 +1,704 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lfm2_moe/modular_lfm2_moe.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_lfm2_moe.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 typing import Optional + +import torch +import torch.nn.functional as F +from torch import nn + +from ... import initialization as init +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, + use_kernelized_func, +) +from ...masking_utils import create_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, 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, can_return_tuple +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.import_utils import is_causal_conv1d_available, is_torchdynamo_compiling +from ...utils.output_capturing import capture_outputs +from .configuration_lfm2_moe import Lfm2MoeConfig + + +if is_causal_conv1d_available(): + from causal_conv1d import causal_conv1d_fn, causal_conv1d_update +else: + causal_conv1d_fn, causal_conv1d_update = None, None + + +@use_kernel_forward_from_hub("RMSNorm") +class Lfm2MoeRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + Lfm2MoeRMSNorm 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 Lfm2MoeRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: Lfm2MoeConfig, 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: Lfm2MoeConfig | 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) + + +class Lfm2MoeMLP(nn.Module): + def __init__(self, config: Lfm2MoeConfig, intermediate_size: int | None = None): + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size + self.w1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.w3 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.w2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + + def forward(self, x): + return self.w2(F.silu(self.w1(x)) * self.w3(x)) + + +@use_experts_implementation +class Lfm2MoeExperts(nn.Module): + """Collection of expert weights stored as 3D tensors.""" + + def __init__(self, config): + super().__init__() + self.num_experts = config.num_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 = F.silu + + 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 Lfm2MoeSparseMoeBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.routed_scaling_factor = config.routed_scaling_factor + self.norm_topk_prob = config.norm_topk_prob + self.use_expert_bias = config.use_expert_bias + + self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) + self.experts = Lfm2MoeExperts(config) + if self.use_expert_bias: + self.register_buffer("expert_bias", torch.zeros(config.num_experts, dtype=torch.float32)) + + def route_tokens_to_experts(self, router_logits): + routing_weights = router_logits.sigmoid() + if self.use_expert_bias: + scores_for_routing = routing_weights + self.expert_bias + _, selected_experts = torch.topk(scores_for_routing, k=self.top_k, dim=-1) + routing_weights = torch.gather(routing_weights, dim=1, index=selected_experts).type_as(router_logits) + else: + routing_weights, selected_experts = torch.topk(routing_weights, k=self.top_k, dim=-1) + + if self.norm_topk_prob: + routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + 1e-6) + routing_weights = routing_weights * self.routed_scaling_factor + return selected_experts, routing_weights + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states_reshaped = hidden_states.view(-1, hidden_dim) + router_logits = self.gate(hidden_states_reshaped) + selected_experts, routing_weights = self.route_tokens_to_experts(router_logits) + final_hidden_states = self.experts(hidden_states_reshaped, selected_experts, routing_weights) + return final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + + +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 Lfm2MoeAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: Lfm2MoeConfig, 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.is_causal = True + self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False) + self.out_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) + self.q_layernorm = Lfm2MoeRMSNorm(self.head_dim, eps=config.norm_eps) + self.k_layernorm = Lfm2MoeRMSNorm(self.head_dim, eps=config.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, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_layernorm(self.q_proj(hidden_states).view(*hidden_shape)).transpose(1, 2) + key_states = self.k_layernorm(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, + scaling=self.scaling, + **kwargs, + ) + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + output = self.out_proj(attn_output) + return output, attn_weights + + +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 + + +kernel_modules = (causal_conv1d_fn, causal_conv1d_update) +is_fast_path_available = all(kernel_modules) + + +class Lfm2MoeShortConv(nn.Module): + def __init__( + self, + config: Lfm2MoeConfig, + layer_idx: int, + ): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.L_cache = config.conv_L_cache + self.bias = config.conv_bias + + self.conv = nn.Conv1d( + in_channels=config.hidden_size, + out_channels=config.hidden_size, + kernel_size=self.L_cache, + groups=config.hidden_size, + bias=self.bias, + padding=self.L_cache - 1, + ) + self.in_proj = nn.Linear(config.hidden_size, 3 * config.hidden_size, bias=self.bias) + self.out_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=self.bias) + + def cuda_kernels_forward( + self, + x: torch.Tensor, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + ): + x = apply_mask_to_padding_states(x, attention_mask) + BCx = self.in_proj(x).transpose(-1, -2) + B, C, x = BCx.chunk(3, dim=-2) + + Bx = B * x + + conv_weights = self.conv.weight.view(self.conv.weight.size(0), self.conv.weight.size(2)) + if past_key_values is not None and past_key_values.has_previous_state(self.layer_idx): + conv_out = causal_conv1d_update( + Bx.squeeze(-1), + past_key_values.layers[self.layer_idx].conv_states, + conv_weights, + self.conv.bias, + None, + ) + conv_out = conv_out.unsqueeze(-1) + else: + if past_key_values is not None: + conv_state = nn.functional.pad(Bx, (self.L_cache - Bx.shape[-1], 0)) + conv_state = past_key_values.update_conv_state(conv_state, self.layer_idx) + + conv_out = causal_conv1d_fn(Bx, conv_weights, self.conv.bias, activation=None) + + y = C * conv_out + y = self.out_proj(y.transpose(-1, -2).contiguous()) + return y + + def slow_forward( + self, + x: torch.Tensor, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + ): + seqlen = x.shape[1] + + x = apply_mask_to_padding_states(x, attention_mask) + BCx = self.in_proj(x).transpose(-1, -2) + B, C, x = BCx.chunk(3, dim=-2) + + Bx = B * x + + if past_key_values is not None and past_key_values.has_previous_state(self.layer_idx): + conv_state = past_key_values.update_conv_state(Bx, self.layer_idx) + conv_out = torch.sum(conv_state.to(Bx.device) * self.conv.weight[:, 0, :], dim=-1) + if self.bias: + conv_out += self.conv.bias + + conv_out = conv_out.unsqueeze(-1) + else: + if past_key_values is not None: + conv_state = nn.functional.pad(Bx, (self.L_cache - Bx.shape[-1], 0)) + conv_state = past_key_values.update_conv_state(conv_state, self.layer_idx) + + conv_out = self.conv(Bx)[..., :seqlen] + + y = C * conv_out + y = y.transpose(-1, -2).contiguous() + y = self.out_proj(y) + return y + + def forward( + self, + hidden_states: torch.Tensor, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + ): + if is_fast_path_available and "cuda" in hidden_states.device.type and not is_torchdynamo_compiling(): + return self.cuda_kernels_forward(hidden_states, past_key_values, attention_mask) + return self.slow_forward(hidden_states, past_key_values, attention_mask) + + +class Lfm2MoeDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: Lfm2MoeConfig, layer_idx: int): + super().__init__() + self.is_attention_layer = config.layer_types[layer_idx] == "full_attention" + + if self.is_attention_layer: + self.self_attn = Lfm2MoeAttention(config, layer_idx) + else: + self.conv = Lfm2MoeShortConv(config, layer_idx) + self.feed_forward = ( + Lfm2MoeMLP(config, intermediate_size=config.intermediate_size) + if layer_idx < config.num_dense_layers + else Lfm2MoeSparseMoeBlock(config) + ) + self.operator_norm = Lfm2MoeRMSNorm(config.hidden_size, eps=config.norm_eps) + self.ffn_norm = Lfm2MoeRMSNorm(config.hidden_size, eps=config.norm_eps) + + 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, + **kwargs, + ) -> torch.Tensor: + residual = hidden_states + if self.is_attention_layer: + hidden_states, _ = self.self_attn( + hidden_states=self.operator_norm(hidden_states), + position_embeddings=position_embeddings, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + **kwargs, + ) + else: + hidden_states = self.conv( + hidden_states=self.operator_norm(hidden_states), + past_key_values=past_key_values, + attention_mask=attention_mask, + ) + hidden_states = hidden_states + residual + hidden_states = hidden_states + self.feed_forward(self.ffn_norm(hidden_states)) + + return hidden_states + + +@auto_docstring +class Lfm2MoePreTrainedModel(PreTrainedModel): + config: Lfm2MoeConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["Lfm2MoeDecoderLayer"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + _can_compile_fullgraph = False # uses a non-compilable cache class + _supports_attention_backend = True + _can_record_outputs = { + "hidden_states": Lfm2MoeDecoderLayer, + "attentions": Lfm2MoeAttention, + } + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + if isinstance(module, Lfm2MoeExperts): + 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) + elif isinstance(module, Lfm2MoeSparseMoeBlock): + if module.use_expert_bias: + init.zeros_(module.expert_bias) + + +@auto_docstring +class Lfm2MoeModel(Lfm2MoePreTrainedModel): + def __init__(self, config: Lfm2MoeConfig): + 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( + [Lfm2MoeDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.gradient_checkpointing = False + self.pos_emb = Lfm2MoeRotaryEmbedding(config) + self.embedding_norm = Lfm2MoeRMSNorm(config.hidden_size, eps=config.norm_eps) + + # 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 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) + + 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, + ) + # Skip masking for decoding stage. We check shape here to be compile-friendly + linear_attention = attention_mask if inputs_embeds.shape[1] != 1 else None + + hidden_states = inputs_embeds + position_embeddings = self.pos_emb(hidden_states, position_ids=position_ids) + + # decoder layers + for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): + layer_mask = causal_mask if self.config.layer_types[i] == "full_attention" else linear_attention + hidden_states = decoder_layer( + hidden_states, + attention_mask=layer_mask, + position_ids=position_ids, + past_key_values=past_key_values, + position_embeddings=position_embeddings, + **kwargs, + ) + + hidden_states = self.embedding_norm(hidden_states) + + return MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +@auto_docstring +class Lfm2MoeForCausalLM(Lfm2MoePreTrainedModel, 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 = Lfm2MoeModel(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, Lfm2MoeForCausalLM + + >>> model = Lfm2MoeForCausalLM.from_pretrained("meta-lfm2_moe/Lfm2Moe-2-7b-hf") + >>> tokenizer = AutoTokenizer.from_pretrained("meta-lfm2_moe/Lfm2Moe-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__ = ["Lfm2MoeForCausalLM", "Lfm2MoeModel", "Lfm2MoePreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/lfm2_moe/modular_lfm2_moe.py b/third_party/transformers/src/transformers/models/lfm2_moe/modular_lfm2_moe.py new file mode 100644 index 0000000000000000000000000000000000000000..a1b2799d2bae88aacb9083c943794b4f9f46c5bc --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2_moe/modular_lfm2_moe.py @@ -0,0 +1,214 @@ +# 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 +import torch.nn.functional as F +from torch import nn + +from ... import initialization as init +from ...cache_utils import Cache, DynamicCache +from ...masking_utils import create_causal_mask +from ...modeling_outputs import MoeModelOutputWithPast +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, logging +from ...utils.import_utils import is_causal_conv1d_available +from ..lfm2.modeling_lfm2 import ( + Lfm2Attention, + Lfm2DecoderLayer, + Lfm2MLP, + Lfm2RotaryEmbedding, + Lfm2ShortConv, +) +from ..llama.modeling_llama import LlamaForCausalLM, LlamaPreTrainedModel, LlamaRMSNorm +from ..mixtral.modeling_mixtral import MixtralModel +from ..qwen2_moe.modeling_qwen2_moe import Qwen2MoeExperts +from .configuration_lfm2_moe import Lfm2MoeConfig + + +if is_causal_conv1d_available(): + from causal_conv1d import causal_conv1d_fn, causal_conv1d_update +else: + causal_conv1d_fn, causal_conv1d_update = None, None + + +kernel_modules = (causal_conv1d_fn, causal_conv1d_update) +is_fast_path_available = all(kernel_modules) + + +logger = logging.get_logger(__name__) + + +class Lfm2MoeRMSNorm(LlamaRMSNorm): + pass + + +class Lfm2MoeRotaryEmbedding(Lfm2RotaryEmbedding): + pass + + +class Lfm2MoeMLP(Lfm2MLP): + def __init__(self, config: Lfm2MoeConfig, intermediate_size: int | None = None): + nn.Module.__init__(self) + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size + self.w1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.w3 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.w2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + + +class Lfm2MoeExperts(Qwen2MoeExperts): + def __init__(self, config): + super().__init__(config) + self.act_fn = F.silu + + +class Lfm2MoeSparseMoeBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.routed_scaling_factor = config.routed_scaling_factor + self.norm_topk_prob = config.norm_topk_prob + self.use_expert_bias = config.use_expert_bias + + self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False) + self.experts = Lfm2MoeExperts(config) + if self.use_expert_bias: + self.register_buffer("expert_bias", torch.zeros(config.num_experts, dtype=torch.float32)) + + def route_tokens_to_experts(self, router_logits): + routing_weights = router_logits.sigmoid() + if self.use_expert_bias: + scores_for_routing = routing_weights + self.expert_bias + _, selected_experts = torch.topk(scores_for_routing, k=self.top_k, dim=-1) + routing_weights = torch.gather(routing_weights, dim=1, index=selected_experts).type_as(router_logits) + else: + routing_weights, selected_experts = torch.topk(routing_weights, k=self.top_k, dim=-1) + + if self.norm_topk_prob: + routing_weights = routing_weights / (routing_weights.sum(dim=-1, keepdim=True) + 1e-6) + routing_weights = routing_weights * self.routed_scaling_factor + return selected_experts, routing_weights + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, sequence_length, hidden_dim = hidden_states.shape + hidden_states_reshaped = hidden_states.view(-1, hidden_dim) + router_logits = self.gate(hidden_states_reshaped) + selected_experts, routing_weights = self.route_tokens_to_experts(router_logits) + final_hidden_states = self.experts(hidden_states_reshaped, selected_experts, routing_weights) + return final_hidden_states.reshape(batch_size, sequence_length, hidden_dim) + + +class Lfm2MoeAttention(Lfm2Attention): + pass + + +class Lfm2MoeShortConv(Lfm2ShortConv): + pass + + +class Lfm2MoeDecoderLayer(Lfm2DecoderLayer): + def __init__(self, config: Lfm2MoeConfig, layer_idx: int): + super().__init__(config, layer_idx) + self.feed_forward = ( + Lfm2MoeMLP(config, intermediate_size=config.intermediate_size) + if layer_idx < config.num_dense_layers + else Lfm2MoeSparseMoeBlock(config) + ) + + +class Lfm2MoePreTrainedModel(LlamaPreTrainedModel): + _can_compile_fullgraph = False # uses a non-compilable cache class + + @torch.no_grad() + def _init_weights(self, module): + PreTrainedModel._init_weights(self, module) + if isinstance(module, Lfm2MoeExperts): + 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) + elif isinstance(module, Lfm2MoeSparseMoeBlock): + if module.use_expert_bias: + init.zeros_(module.expert_bias) + + +class Lfm2MoeModel(MixtralModel): + def __init__(self, config: Lfm2MoeConfig): + super().__init__(config) + self.pos_emb = Lfm2MoeRotaryEmbedding(config) + self.embedding_norm = Lfm2MoeRMSNorm(config.hidden_size, eps=config.norm_eps) + del self.norm + del self.rotary_emb + + 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 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) + + 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, + ) + # Skip masking for decoding stage. We check shape here to be compile-friendly + linear_attention = attention_mask if inputs_embeds.shape[1] != 1 else None + + hidden_states = inputs_embeds + position_embeddings = self.pos_emb(hidden_states, position_ids=position_ids) + + # decoder layers + for i, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): + layer_mask = causal_mask if self.config.layer_types[i] == "full_attention" else linear_attention + hidden_states = decoder_layer( + hidden_states, + attention_mask=layer_mask, + position_ids=position_ids, + past_key_values=past_key_values, + position_embeddings=position_embeddings, + **kwargs, + ) + + hidden_states = self.embedding_norm(hidden_states) + + return MoeModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +class Lfm2MoeForCausalLM(LlamaForCausalLM): + pass + + +__all__ = ["Lfm2MoeForCausalLM", "Lfm2MoeModel", "Lfm2MoePreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/lfm2_vl/__init__.py b/third_party/transformers/src/transformers/models/lfm2_vl/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c551d83516d8e7a93fa26c1ce6f1524ce40bc9cd --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2_vl/__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_lfm2_vl import * + from .image_processing_lfm2_vl import * + from .modeling_lfm2_vl import * + from .processing_lfm2_vl 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/lfm2_vl/configuration_lfm2_vl.py b/third_party/transformers/src/transformers/models/lfm2_vl/configuration_lfm2_vl.py new file mode 100644 index 0000000000000000000000000000000000000000..40ef518e41f036f2b3e37842fec2b01e7621f967 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2_vl/configuration_lfm2_vl.py @@ -0,0 +1,63 @@ +# 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. +"""PyTorch LFM2-VL model.""" + +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="LiquidAI/LFM2-VL-1.6B") +@strict +class Lfm2VlConfig(PreTrainedConfig): + r""" + projector_use_layernorm (`bool`, *optional*, defaults to `True`): + Whether to use layernorm in the multimodal projector. + downsample_factor (`int`, *optional*, defaults to 2): + The downsample_factor factor of the vision backbone. + """ + + model_type = "lfm2_vl" + sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig} + + vision_config: dict | PreTrainedConfig | None = None + text_config: dict | PreTrainedConfig | None = None + image_token_id: int = 396 + projector_hidden_act: str = "gelu" + projector_hidden_size: int = 2560 + projector_bias: bool = True + projector_use_layernorm: bool = True + downsample_factor: int = 2 + tie_word_embeddings: bool = True + + def __post_init__(self, **kwargs): + if isinstance(self.vision_config, dict): + self.vision_config["model_type"] = self.vision_config.get("model_type", "siglip2_vision_model") + self.vision_config = CONFIG_MAPPING[self.vision_config["model_type"]](**self.vision_config) + elif self.vision_config is None: + self.vision_config = CONFIG_MAPPING["siglip2_vision_model"]() + + if isinstance(self.text_config, dict): + self.text_config["model_type"] = self.text_config.get("model_type", "lfm2") + self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config) + elif self.text_config is None: + self.text_config = CONFIG_MAPPING["lfm2"]() + + self.tie_word_embeddings = kwargs.pop("tie_embedding", self.tie_word_embeddings) + super().__post_init__(**kwargs) + + +__all__ = ["Lfm2VlConfig"] diff --git a/third_party/transformers/src/transformers/models/lfm2_vl/image_processing_lfm2_vl.py b/third_party/transformers/src/transformers/models/lfm2_vl/image_processing_lfm2_vl.py new file mode 100644 index 0000000000000000000000000000000000000000..70cf5c06f02c69eb2733410ecebc6ad6b6932ac0 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2_vl/image_processing_lfm2_vl.py @@ -0,0 +1,561 @@ +# 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 math +from functools import lru_cache + +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, split_to_tiles +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + PILImageResampling, + SizeDict, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import ( + TensorType, + auto_docstring, + logging, +) + + +logger = logging.get_logger(__name__) + + +def round_by_factor(number: float, factor: int) -> int: + """Returns the closest integer to 'number' that is divisible by 'factor'.""" + return round(number / factor) * factor + + +def find_closest_aspect_ratio( + aspect_ratio: float, + target_ratios: list[tuple[int, int]], + width: int, + height: int, + image_size: int, +) -> tuple[int, int]: + """Find the closest aspect ratio from target_ratios to match the input aspect ratio. + + Args: + aspect_ratio: The aspect ratio to match (width/height). + target_ratios: List of possible aspect ratios as tuples of (width, height) integers. + width: Original image width in pixels. + height: Original image height in pixels. + image_size: Base size for calculating target area. + + Returns: + tuple[int, int]: The best matching ratio as (width, height) integers. + """ + best_ratio_diff = float("inf") + best_ratio = (1, 1) + area = width * height + + for ratio in target_ratios: + target_aspect_ratio = ratio[0] / ratio[1] + ratio_diff = abs(aspect_ratio - target_aspect_ratio) + + # update best ratio if we found a closer match + if ratio_diff < best_ratio_diff: + best_ratio_diff = ratio_diff + best_ratio = ratio + # if equally close, prefer the ratio that better matches the original image area + elif ratio_diff == best_ratio_diff: + target_area = image_size * image_size * ratio[0] * ratio[1] + if area > 0.5 * target_area: + best_ratio = ratio + + return best_ratio + + +# copied from Siglip2ImageProcessor +@lru_cache(maxsize=256) +def get_image_size_for_max_num_patches( + image_height: int, image_width: int, patch_size: int, max_num_patches: int, eps: float = 1e-5 +) -> tuple[int, int]: + """ + Determine image size based on max number of patches, ensure dimensions are divisible by patch size and image is at least 1 patch. + + Args: + image_height (`int`): + Original image height. + image_width (`int`): + Original image width. + patch_size (`int`): + Patch size for processing. + max_num_patches (`int`): + Maximum number of patches. + eps (`float`): + Small threshold for binary search. + + Returns: + Tuple: (target_height, target_width) + """ + + def get_scaled_image_size(scale: float, size: int, patch_size: int) -> int: + scaled_size = size * scale + scaled_size = math.ceil(scaled_size / patch_size) * patch_size # make divisible by patch_size + scaled_size = max(patch_size, scaled_size) # ensure at least 1 patch + return int(scaled_size) + + # Binary search for optimal scale + scale_min, scale_max = eps / 10, 100.0 + while (scale_max - scale_min) >= eps: + scale = (scale_min + scale_max) / 2 + target_height = get_scaled_image_size(scale, image_height, patch_size) + target_width = get_scaled_image_size(scale, image_width, patch_size) + num_patches = (target_height / patch_size) * (target_width / patch_size) + + if num_patches <= max_num_patches: + scale_min = scale + else: + scale_max = scale + + scale = scale_min + target_height = get_scaled_image_size(scale, image_height, patch_size) + target_width = get_scaled_image_size(scale, image_width, patch_size) + return target_height, target_width + + +def convert_image_to_patches(images: "torch.Tensor", patch_size: int) -> "torch.Tensor": + """ + Convert 3D array image of shape (image_height, image_width, num_channels) into 2D array of patches of shape + (num_patches_height * num_patches_width, patch_size * patch_size * num_channels). + """ + batch_size, num_channels, image_height, image_width = images.shape + num_patches_height = image_height // patch_size + num_patches_width = image_width // patch_size + patched_image = images.reshape( + batch_size, num_channels, num_patches_height, patch_size, num_patches_width, patch_size + ) + patched_image = patched_image.permute(0, 2, 4, 3, 5, 1) + patched_image = patched_image.reshape(batch_size, num_patches_height * num_patches_width, -1) + return patched_image + + +def pad_along_first_dim( + images: "torch.Tensor", target_length: int, pad_value: int = 0 +) -> tuple["torch.Tensor", "torch.Tensor"]: + """ + Pad the array along the first dimension. + """ + current_length = images.shape[1] + padding_length = target_length - current_length + pixel_mask = torch.ones((target_length,), dtype=torch.int32) + if padding_length > 0: + paddings = (0, 0, 0, padding_length, 0, 0) + images = torch.nn.functional.pad(images, paddings, mode="constant", value=pad_value) + pixel_mask[-padding_length:] = 0 + return images, pixel_mask + + +class Lfm2VlImageProcessorKwargs(ImagesKwargs, total=False): + """ + downsample_factor (`int`, *optional*, defaults to `2`): + The downsampling factor for images used when resizing the image. + do_image_splitting (`bool`, *optional*, defaults to `True`): + Whether to split large images into a grid of smaller tiles. When enabled, images exceeding the maximum token + limit are divided into multiple tiles based on `min_tiles` and `max_tiles` constraints. + min_tiles (`int`, *optional*, defaults to `2`): + Minimum number of tiles (width × height) to use when splitting an image into a grid. The grid configuration + is chosen to maintain the original aspect ratio while staying within the `min_tiles` and `max_tiles` range. + max_tiles (`int`, *optional*, defaults to `10`): + Maximum number of tiles (width × height) to use when splitting an image into a grid. The grid configuration + is chosen to maintain the original aspect ratio while staying within the `min_tiles` and `max_tiles` range. + use_thumbnail (`bool`, *optional*, defaults to `True`): + Whether to include a thumbnail version of the image when splitting into tiles. The thumbnail provides a + low-resolution overview of the entire image and is added as an additional patch when the grid has more than + one tile. + min_image_tokens (`int`, *optional*, defaults to `64`): + Minimum number of image tokens (patches) to generate for an image. Images smaller than this threshold will + be upscaled to meet the minimum token requirement. + max_image_tokens (`int`, *optional*, defaults to `256`): + Maximum number of image tokens (patches) allowed for a single image. Images exceeding this limit will be + split into multiple tiles or downscaled accordingly. + encoder_patch_size (`int`, *optional*, defaults to `16`): + The patch size used by the vision encoder. Images are divided into patches of this size, and both height + and width must be divisible by this value (after accounting for the downsampling factor). + tile_size (`int`, *optional*, defaults to `512`): + The size of each tile when splitting large images into a grid. Each tile will be resized to this dimension + before being processed into patches. + max_pixels_tolerance (`float`, *optional*, defaults to `2.0`): + Tolerance factor for determining if an image is too large. An image is considered too large if its pixel + count exceeds `max_image_tokens * encoder_patch_size^2 * downsample_factor^2 * max_pixels_tolerance`. + return_row_col_info (`bool`, *optional*, defaults to `False`): + Whether to return row and column information for each image in the batch. When enabled, the output includes + `image_rows`, `image_cols`, and `image_sizes` fields indicating the grid layout and dimensions of processed images. + """ + + downsample_factor: int + do_image_splitting: bool + min_tiles: int + max_tiles: int + use_thumbnail: bool + min_image_tokens: int + max_image_tokens: int + encoder_patch_size: int + tile_size: int + max_pixels_tolerance: float + do_pad: bool + return_row_col_info: bool + + +@auto_docstring +class Lfm2VlImageProcessor(TorchvisionBackend): + downsample_factor = 2 + do_image_splitting = True + min_tiles = 2 + max_tiles = 10 + use_thumbnail = True + min_image_tokens = 64 + max_image_tokens = 256 + encoder_patch_size = 16 + tile_size = 512 + max_pixels_tolerance = 2.0 + do_resize = True + size = {"height": 512, "width": 512} + resample = PILImageResampling.BILINEAR + do_rescale = True + rescale_factor = 1 / 255 + do_normalize = True + do_pad = True + return_row_col_info = False + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + valid_kwargs = Lfm2VlImageProcessorKwargs + model_input_names = ["pixel_values", "pixel_attention_mask", "spatial_shapes"] + + def __init__(self, **kwargs: Unpack[Lfm2VlImageProcessorKwargs]): + super().__init__(**kwargs) + + max_thumbnail_image_patches = self.max_image_tokens * self.downsample_factor**2 + tile_size_patches = (self.tile_size // self.encoder_patch_size) ** 2 if self.do_image_splitting else 0 + self.max_num_patches = max( + max_thumbnail_image_patches, + tile_size_patches, + ) + + @lru_cache(maxsize=256) + def _target_ratios(self, min_tiles: int, max_tiles: int) -> list[tuple[int, int]]: + ratios = [ + (w, h) + for n in range(min_tiles, max_tiles + 1) + for w in range(1, n + 1) + for h in range(1, n + 1) + if min_tiles <= w * h <= max_tiles + ] + return sorted(set(ratios), key=lambda x: x[0] * x[1]) + + def _get_grid_layout( + self, + height: int, + width: int, + min_tiles: int, + max_tiles: int, + tile_size: int, + ) -> tuple[int, int]: + aspect_ratio = width / height + target_ratios = self._target_ratios(min_tiles, max_tiles) + + # find best matching grid configuration + grid_width, grid_height = find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, tile_size) + + target_width = tile_size * grid_width + target_height = tile_size * grid_height + total_patches = grid_width * grid_height + + return grid_width, grid_height, target_width, target_height, total_patches + + def crop_image_to_patches( + self, + image: "torch.Tensor", + min_tiles: int, + max_tiles: int, + tile_size: int, + use_thumbnail: bool, + thumbnail_size: tuple[int], + resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None, + antialias: bool = True, + **kwargs, + ) -> "torch.Tensor": + """ + Processes a high resolution image into patches. + This method splits a high resolution image into a grid of smaller patches while trying to maintain + the original aspect ratio. It finds the optimal grid configuration within the specified tile constraints. + """ + batch_size, num_channels, height, width = image.shape + grid_width, grid_height, target_width, target_height, total_patches = self._get_grid_layout( + height, width, min_tiles=min_tiles, max_tiles=max_tiles, tile_size=tile_size + ) + + resized_image = super().resize( + image, SizeDict(height=target_height, width=target_width), resample=resample, antialias=antialias + ) + + # split the image into patches + processed_images = split_to_tiles(resized_image, num_tiles_height=grid_height, num_tiles_width=grid_width) + + # Re-order processed images to a nested image structure, so it can be reordered back correctly + # Note that the images can't be stacked because the thumbnail image is of bigger size than patches + # Each image in sublist will be of shape (1, C, H, W) + processed_images = list(processed_images) + + if use_thumbnail and grid_width * grid_height != 1: + total_patches += 1 + thumbnail_image = super().resize( + image, + SizeDict(height=thumbnail_size[0], width=thumbnail_size[1]), + resample=resample, + antialias=antialias, + ) + for i in range(batch_size): + processed_images[i] = list(processed_images[i]) + list(thumbnail_image[i][None, ...]) + + return processed_images, grid_width, grid_height + + # Adapted from Qwen-VL with minor differences + def smart_resize( + self, + height: int, + width: int, + downsample_factor: int, + min_image_tokens: int, + max_image_tokens: int, + encoder_patch_size: int, + ) -> tuple[int, int]: + """ + Rescales the image so that the following conditions are met: + 1. Both dimensions (height and width) are divisible by 'encoder_patch_size' * 'downsample_factor'. + This ensures no padding is needed in the downsampling step. + 2. The total number of pixels is within the range ['smart_resize_min_pixels', 'smart_resize_max_pixels']. + 3. The aspect ratio of the image is maintained as closely as possible. + """ + total_factor = encoder_patch_size * downsample_factor + smart_resize_min_pixels = min_image_tokens * encoder_patch_size**2 * downsample_factor**2 + smart_resize_max_pixels = max_image_tokens * encoder_patch_size**2 * downsample_factor**2 + + h_bar = max(total_factor, round_by_factor(height, total_factor)) + w_bar = max(total_factor, round_by_factor(width, total_factor)) + + if h_bar * w_bar > smart_resize_max_pixels: + beta = math.sqrt((height * width) / smart_resize_max_pixels) + math.floor(height / beta / total_factor) * total_factor + h_bar = max(total_factor, math.floor(height / beta / total_factor) * total_factor) + w_bar = max(total_factor, math.floor(width / beta / total_factor) * total_factor) + elif h_bar * w_bar < smart_resize_min_pixels: + beta = math.sqrt(smart_resize_min_pixels / (height * width)) + h_bar = math.ceil(height * beta / total_factor) * total_factor + w_bar = math.ceil(width * beta / total_factor) * total_factor + + return w_bar, h_bar + + def _is_image_too_large( + self, + height: int, + width: int, + max_image_tokens: int, + encoder_patch_size: int, + downsample_factor: int, + max_pixels_tolerance: float, + ) -> bool: + """Check if the image is too large to be processed as one tile.""" + total_factor = encoder_patch_size * downsample_factor + + h_bar = max(encoder_patch_size, round_by_factor(height, total_factor)) + w_bar = max(encoder_patch_size, round_by_factor(width, total_factor)) + return h_bar * w_bar > max_image_tokens * encoder_patch_size**2 * downsample_factor**2 * max_pixels_tolerance + + def resize_and_split( + self, + images: "torch.Tensor", + downsample_factor: int, + min_tiles: int, + max_tiles: int, + use_thumbnail: bool, + min_image_tokens: int, + max_image_tokens: int, + encoder_patch_size: int, + tile_size: int, + max_pixels_tolerance: float, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + ) -> tuple[list["torch.Tensor"], list[list[int]], list[list[int]], list[list[tuple[int, int]]]]: + batch_size, _, height, width = images.shape + do_image_splitting = not min_tiles == max_tiles == 1 + is_image_large = self._is_image_too_large( + height=height, + width=width, + max_image_tokens=max_image_tokens, + encoder_patch_size=encoder_patch_size, + downsample_factor=downsample_factor, + max_pixels_tolerance=max_pixels_tolerance, + ) + + new_width, new_height = self.smart_resize( + height=height, + width=width, + downsample_factor=downsample_factor, + min_image_tokens=min_image_tokens, + max_image_tokens=max_image_tokens, + encoder_patch_size=encoder_patch_size, + ) + + # Big image will be cropped into patches and small images are just resized + if is_image_large and do_image_splitting: + images, num_cols, num_rows = self.crop_image_to_patches( + images, + min_tiles=min_tiles, + max_tiles=max_tiles, + tile_size=tile_size, + thumbnail_size=(new_height, new_width), + use_thumbnail=use_thumbnail, + resample=resample, + ) + else: + num_rows = num_cols = 1 + images = super().resize(images, SizeDict(height=new_height, width=new_width), resample=resample) + # Make a list and treat it as single crop per image so it can be re-grouped back correctly + images = [[image] for image in images] + + num_rows = [num_rows] * batch_size + num_cols = [num_cols] * batch_size + image_sizes = [[new_height, new_width]] * batch_size + return images, num_rows, num_cols, image_sizes + + def _preprocess( + self, + images: list["torch.Tensor"], + do_resize: bool, + resample: "PILImageResampling | tvF.InterpolationMode | int | None", + do_rescale: bool, + rescale_factor: float, + do_normalize: bool, + image_mean: float | list[float], + image_std: float | list[float], + downsample_factor: int, + do_image_splitting: bool, + min_tiles: int, + max_tiles: int, + use_thumbnail: bool, + min_image_tokens: int, + max_image_tokens: int, + encoder_patch_size: int, + tile_size: int, + max_pixels_tolerance: float, + do_pad: bool, + return_row_col_info: bool, + return_tensors: str | TensorType | None, + disable_grouping: bool | None, + **kwargs, + ) -> BatchFeature: + if not do_image_splitting: + min_tiles = 1 + max_tiles = 1 + logger.debug( + "Image splitting is disabled, setting min_tiles and max_tiles to 1. Set do_image_splitting=True to enable splitting." + ) + + if do_image_splitting and min_tiles > max_tiles: + raise ValueError("min_tiles must be less than or equal to max_tiles") + + max_thumbnail_image_patches = max_image_tokens * downsample_factor**2 + tile_size_patches = (tile_size // encoder_patch_size) ** 2 if do_image_splitting else 0 + max_num_patches = max( + max_thumbnail_image_patches, + tile_size_patches, + ) + + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + resized_image_sizes = {} + rows_grouped, cols_grouped = {}, {} + for shape, stacked_images in grouped_images.items(): + num_rows = [1] * stacked_images.shape[0] + num_cols = [1] * stacked_images.shape[0] + height, width = stacked_images.shape[-2:] + image_sizes = [[height, width]] * stacked_images.shape[0] + do_resize = True + + if do_resize: + stacked_images, num_rows, num_cols, image_sizes = self.resize_and_split( + stacked_images, + downsample_factor=downsample_factor, + min_tiles=min_tiles, + max_tiles=max_tiles, + use_thumbnail=use_thumbnail, + min_image_tokens=min_image_tokens, + max_image_tokens=max_image_tokens, + encoder_patch_size=encoder_patch_size, + tile_size=tile_size, + max_pixels_tolerance=max_pixels_tolerance, + resample=resample, + ) + + rows_grouped[shape] = num_rows + cols_grouped[shape] = num_cols + resized_image_sizes[shape] = image_sizes + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index) + batch_rows = reorder_images(rows_grouped, grouped_images_index) + batch_cols = reorder_images(cols_grouped, grouped_images_index) + resized_image_sizes = reorder_images(resized_image_sizes, grouped_images_index) + + grouped_images, grouped_images_index = group_images_by_shape( + resized_images, disable_grouping=disable_grouping, is_nested=True + ) + + processed_images_grouped = {} + processed_masks, processed_spatial_shapes = {}, {} + for shape, stacked_images in grouped_images.items(): + # Fused rescale and normalize + stacked_images = self.rescale_and_normalize( + stacked_images, do_rescale, rescale_factor, do_normalize, image_mean, image_std + ) + batch_size, *_, height, width = stacked_images.shape + num_patches_height = height // encoder_patch_size + num_patches_width = width // encoder_patch_size + + stacked_images = convert_image_to_patches(stacked_images, encoder_patch_size) + processed_spatial_shapes[shape] = [[num_patches_height, num_patches_width]] * batch_size + + if do_pad: + stacked_images, pixel_mask = pad_along_first_dim(stacked_images, max_num_patches) + processed_masks[shape] = [pixel_mask] * batch_size + + processed_images_grouped[shape] = stacked_images + + processed_images = reorder_images(processed_images_grouped, grouped_images_index, is_nested=True) + data = {"pixel_values": torch.cat([torch.stack(images) for images in processed_images])} + + if do_pad: + processed_masks = reorder_images(processed_masks, grouped_images_index, is_nested=True) + processed_spatial_shapes = reorder_images(processed_spatial_shapes, grouped_images_index, is_nested=True) + processed_masks = torch.cat([torch.stack(masks) for masks in processed_masks]) + processed_spatial_shapes = torch.cat( + [torch.tensor(spatial_shape) for spatial_shape in processed_spatial_shapes] + ) + data.update({"pixel_attention_mask": processed_masks, "spatial_shapes": processed_spatial_shapes}) + + if return_row_col_info: + data["image_rows"] = batch_rows + data["image_cols"] = batch_cols + data["image_sizes"] = resized_image_sizes + + encoding = BatchFeature(data=data, tensor_type=return_tensors) + return encoding + + +__all__ = ["Lfm2VlImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/lfm2_vl/modeling_lfm2_vl.py b/third_party/transformers/src/transformers/models/lfm2_vl/modeling_lfm2_vl.py new file mode 100644 index 0000000000000000000000000000000000000000..b66ba44bef3f4432b2f89f612eda0055558ec30d --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2_vl/modeling_lfm2_vl.py @@ -0,0 +1,482 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/lfm2_vl/modular_lfm2_vl.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_lfm2_vl.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# 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. + +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, can_return_tuple, torch_compilable_check +from ..auto import AutoModel +from .configuration_lfm2_vl import Lfm2VlConfig + + +class Lfm2VlMultiModalProjector(nn.Module): + def __init__(self, config: Lfm2VlConfig): + super().__init__() + in_channels = config.vision_config.hidden_size * (config.downsample_factor**2) + self.factor = config.downsample_factor + self.use_layer_norm = config.projector_use_layernorm + self.layer_norm = nn.LayerNorm(in_channels) if config.projector_use_layernorm else None + self.linear_1 = nn.Linear( + in_channels, + config.projector_hidden_size, + bias=config.projector_bias, + ) + self.act = ACT2FN[config.projector_hidden_act] + self.linear_2 = nn.Linear( + config.projector_hidden_size, + config.text_config.hidden_size, + bias=config.projector_bias, + ) + + def forward(self, image_features: torch.Tensor): + image_features = self.pixel_unshuffle(image_features) + if self.use_layer_norm: + image_features = self.layer_norm(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 + + def pixel_unshuffle(self, hidden_states: torch.Tensor): + batch_size, width, height, channels = hidden_states.size() + hidden_states = hidden_states.reshape(batch_size, width, height // self.factor, channels * self.factor) + hidden_states = hidden_states.permute(0, 2, 1, 3) + hidden_states = hidden_states.reshape( + batch_size, height // self.factor, width // self.factor, channels * self.factor**2 + ) + hidden_states = hidden_states.permute(0, 2, 1, 3) + return hidden_states + + +@auto_docstring +class Lfm2VlPreTrainedModel(PreTrainedModel): + config: Lfm2VlConfig + 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 = False + _supports_flex_attn = True + _supports_attention_backend = True + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Lfm2Vl causal language model (or autoregressive) outputs. + """ +) +class Lfm2VlCausalLMOutputWithPast(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 + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Lfm2Vl outputs, with hidden states and attentions. + """ +) +class Lfm2VlModelOutputWithPast(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 Lfm2Vl model which consists of a vision backbone and a language model, without a language modeling head. + """ +) +class Lfm2VlModel(Lfm2VlPreTrainedModel): + def __init__(self, config: Lfm2VlConfig): + super().__init__(config) + self.vision_tower = AutoModel.from_config(config.vision_config) + + self.multi_modal_projector = Lfm2VlMultiModalProjector(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) + + @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, + spatial_shapes: torch.Tensor, + pixel_attention_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`): + The spatial shapes of the input images. + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`): + The pixel attention mask of the input images. + """ + image_outputs = self.vision_tower( + pixel_values=pixel_values, + spatial_shapes=spatial_shapes, + pixel_attention_mask=pixel_attention_mask, + return_dict=True, + **kwargs, + ) + last_hidden_state = image_outputs.last_hidden_state + + img_feature_lengths = pixel_attention_mask.sum(dim=1) + image_features = [] + + for img_idx in range(last_hidden_state.size(0)): + feature = last_hidden_state[img_idx] + # unpad the image representation + feature = feature[: img_feature_lengths[img_idx], :].unsqueeze(0) + + # reshape to original height and width + feature_org_h, feature_org_w = spatial_shapes[img_idx] + feature = feature.reshape(1, feature_org_h, feature_org_w, -1) + + # project the image representation + img_embedding = self.multi_modal_projector(feature) + + # flatten here to handle variable length in naflex + img_embedding = img_embedding.reshape(-1, img_embedding.size(-1)) + image_features.append(img_embedding) + + image_outputs.pooler_output = 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() + special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + n_image_features = image_features.shape[0] + 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, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + spatial_shapes: torch.Tensor | None = None, + pixel_attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Lfm2VlModelOutputWithPast: + r""" + spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`, *optional*): + The spatial shapes of the input images. + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`, *optional*): + The pixel attention mask of the input images. + """ + + 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, + spatial_shapes=spatial_shapes, + pixel_attention_mask=pixel_attention_mask, + 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=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, + use_cache=use_cache, + **kwargs, + ) + + return Lfm2VlModelOutputWithPast( + 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, + ) + + +@auto_docstring( + custom_intro=""" + The LFM2_VL model which consists of a vision backbone and a language model. + """ +) +class Lfm2VlForConditionalGeneration(Lfm2VlPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"} + + def __init__(self, config: Lfm2VlConfig): + super().__init__(config) + self.model = Lfm2VlModel(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, + spatial_shapes: torch.Tensor, + pixel_attention_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`): + The spatial shapes of the input images. + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`): + The pixel attention mask of the input images. + """ + return self.model.get_image_features( + pixel_values=pixel_values, + spatial_shapes=spatial_shapes, + pixel_attention_mask=pixel_attention_mask, + **kwargs, + ) + + @can_return_tuple + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + spatial_shapes: torch.Tensor | None = None, + pixel_attention_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, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Lfm2VlCausalLMOutputWithPast: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, channels, height, width)`, *optional*): + The input image tensors. + spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`, *optional*): + The spatial shapes of the input images. + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`, *optional*): + The pixel attention mask of the input 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 PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> from transformers import AutoProcessor, AutoModelForImageTextToText + >>> from transformers.image_utils import load_image + + >>> model = AutoModelForImageTextToText.from_pretrained( + ... "LiquidAI/LFM2-VL-1.6B", + ... ) + >>> processor = AutoProcessor.from_pretrained( + ... "LiquidAI/LFM2-VL-1.6B", + ... ) + + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> image = load_image(url) + + >>> conversation = [ + ... { + ... "role": "user", + ... "content": [ + ... {"type": "image", "image": image}, + ... {"type": "text", "text": "What is in this image?"}, + ... ], + ... }, + ... ] + + >>> inputs = processor.apply_chat_template( + ... conversation, + ... add_generation_prompt=True, + ... tokenize=True, + ... return_dict=True, + ... return_tensors="pt" + ... ) + + >>> # Generate + >>> outputs = model.generate(**inputs, max_new_tokens=45) + >>> processor.batch_decode(outputs, skip_special_tokens=True)[0] + 'This image depicts a vibrant street scene in what appears to be a Chinatown or similar cultural area. The focal point is a large red stop sign with white lettering, mounted on a pole.' + ```""" + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + spatial_shapes=spatial_shapes, + pixel_attention_mask=pixel_attention_mask, + 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 Lfm2VlCausalLMOutputWithPast( + 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__ = ["Lfm2VlForConditionalGeneration", "Lfm2VlPreTrainedModel", "Lfm2VlModel"] diff --git a/third_party/transformers/src/transformers/models/lfm2_vl/modular_lfm2_vl.py b/third_party/transformers/src/transformers/models/lfm2_vl/modular_lfm2_vl.py new file mode 100644 index 0000000000000000000000000000000000000000..4cf94132367cf62c5be5b9f6cc55f9a81e9e99a5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2_vl/modular_lfm2_vl.py @@ -0,0 +1,359 @@ +# 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. +"""PyTorch Lfm2-VL model.""" + +import torch +from torch import nn + +from ...activations import ACT2FN +from ...cache_utils import Cache +from ...modeling_outputs import BaseModelOutputWithPooling +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check +from ..llava.modeling_llava import ( + LlavaCausalLMOutputWithPast, + LlavaForConditionalGeneration, + LlavaModel, + LlavaModelOutputWithPast, + LlavaPreTrainedModel, +) +from .configuration_lfm2_vl import Lfm2VlConfig + + +logger = logging.get_logger(__name__) + + +class Lfm2VlMultiModalProjector(nn.Module): + def __init__(self, config: Lfm2VlConfig): + super().__init__() + in_channels = config.vision_config.hidden_size * (config.downsample_factor**2) + self.factor = config.downsample_factor + self.use_layer_norm = config.projector_use_layernorm + self.layer_norm = nn.LayerNorm(in_channels) if config.projector_use_layernorm else None + self.linear_1 = nn.Linear( + in_channels, + config.projector_hidden_size, + bias=config.projector_bias, + ) + self.act = ACT2FN[config.projector_hidden_act] + self.linear_2 = nn.Linear( + config.projector_hidden_size, + config.text_config.hidden_size, + bias=config.projector_bias, + ) + + def forward(self, image_features: torch.Tensor): + image_features = self.pixel_unshuffle(image_features) + if self.use_layer_norm: + image_features = self.layer_norm(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 + + def pixel_unshuffle(self, hidden_states: torch.Tensor): + batch_size, width, height, channels = hidden_states.size() + hidden_states = hidden_states.reshape(batch_size, width, height // self.factor, channels * self.factor) + hidden_states = hidden_states.permute(0, 2, 1, 3) + hidden_states = hidden_states.reshape( + batch_size, height // self.factor, width // self.factor, channels * self.factor**2 + ) + hidden_states = hidden_states.permute(0, 2, 1, 3) + return hidden_states + + +class Lfm2VlPreTrainedModel(LlavaPreTrainedModel): + _can_compile_fullgraph = False + base_model_prefix = "model" + + +class Lfm2VlCausalLMOutputWithPast(LlavaCausalLMOutputWithPast): + pass + + +class Lfm2VlModelOutputWithPast(LlavaModelOutputWithPast): + pass + + +class Lfm2VlModel(LlavaModel): + def __init__(self, config: Lfm2VlConfig): + super().__init__(config) + + @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, + spatial_shapes: torch.Tensor, + pixel_attention_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`): + The spatial shapes of the input images. + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`): + The pixel attention mask of the input images. + """ + image_outputs = self.vision_tower( + pixel_values=pixel_values, + spatial_shapes=spatial_shapes, + pixel_attention_mask=pixel_attention_mask, + return_dict=True, + **kwargs, + ) + last_hidden_state = image_outputs.last_hidden_state + + img_feature_lengths = pixel_attention_mask.sum(dim=1) + image_features = [] + + for img_idx in range(last_hidden_state.size(0)): + feature = last_hidden_state[img_idx] + # unpad the image representation + feature = feature[: img_feature_lengths[img_idx], :].unsqueeze(0) + + # reshape to original height and width + feature_org_h, feature_org_w = spatial_shapes[img_idx] + feature = feature.reshape(1, feature_org_h, feature_org_w, -1) + + # project the image representation + img_embedding = self.multi_modal_projector(feature) + + # flatten here to handle variable length in naflex + img_embedding = img_embedding.reshape(-1, img_embedding.size(-1)) + image_features.append(img_embedding) + + image_outputs.pooler_output = 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() + special_image_mask = special_image_mask.unsqueeze(-1).expand_as(inputs_embeds).to(inputs_embeds.device) + n_image_features = image_features.shape[0] + 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, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + spatial_shapes: torch.Tensor | None = None, + pixel_attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Lfm2VlModelOutputWithPast: + r""" + spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`, *optional*): + The spatial shapes of the input images. + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`, *optional*): + The pixel attention mask of the input images. + """ + + 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, + spatial_shapes=spatial_shapes, + pixel_attention_mask=pixel_attention_mask, + 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=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, + use_cache=use_cache, + **kwargs, + ) + + return Lfm2VlModelOutputWithPast( + 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 Lfm2VlForConditionalGeneration(LlavaForConditionalGeneration): + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + spatial_shapes: torch.Tensor, + pixel_attention_mask: torch.Tensor, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPooling: + r""" + pixel_values (`torch.FloatTensor]` of shape `(batch_size, channels, height, width)`): + The tensors corresponding to the input images. + spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`): + The spatial shapes of the input images. + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`): + The pixel attention mask of the input images. + """ + return self.model.get_image_features( + pixel_values=pixel_values, + spatial_shapes=spatial_shapes, + pixel_attention_mask=pixel_attention_mask, + **kwargs, + ) + + @can_return_tuple + def forward( + self, + input_ids: torch.LongTensor | None = None, + pixel_values: torch.FloatTensor | None = None, + spatial_shapes: torch.Tensor | None = None, + pixel_attention_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, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | Lfm2VlCausalLMOutputWithPast: + r""" + pixel_values (`torch.FloatTensor` of shape `(batch_size, channels, height, width)`, *optional*): + The input image tensors. + spatial_shapes (`torch.Tensor` of shape `(batch_size, 2)`, *optional*): + The spatial shapes of the input images. + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, height, width)`, *optional*): + The pixel attention mask of the input 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 PIL import Image + >>> import httpx + >>> from io import BytesIO + >>> from transformers import AutoProcessor, AutoModelForImageTextToText + >>> from transformers.image_utils import load_image + + >>> model = AutoModelForImageTextToText.from_pretrained( + ... "LiquidAI/LFM2-VL-1.6B", + ... ) + >>> processor = AutoProcessor.from_pretrained( + ... "LiquidAI/LFM2-VL-1.6B", + ... ) + + >>> url = "https://www.ilankelman.org/stopsigns/australia.jpg" + >>> image = load_image(url) + + >>> conversation = [ + ... { + ... "role": "user", + ... "content": [ + ... {"type": "image", "image": image}, + ... {"type": "text", "text": "What is in this image?"}, + ... ], + ... }, + ... ] + + >>> inputs = processor.apply_chat_template( + ... conversation, + ... add_generation_prompt=True, + ... tokenize=True, + ... return_dict=True, + ... return_tensors="pt" + ... ) + + >>> # Generate + >>> outputs = model.generate(**inputs, max_new_tokens=45) + >>> processor.batch_decode(outputs, skip_special_tokens=True)[0] + 'This image depicts a vibrant street scene in what appears to be a Chinatown or similar cultural area. The focal point is a large red stop sign with white lettering, mounted on a pole.' + ```""" + outputs = self.model( + input_ids=input_ids, + pixel_values=pixel_values, + spatial_shapes=spatial_shapes, + pixel_attention_mask=pixel_attention_mask, + 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 Lfm2VlCausalLMOutputWithPast( + 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__ = ["Lfm2VlForConditionalGeneration", "Lfm2VlPreTrainedModel", "Lfm2VlModel"] diff --git a/third_party/transformers/src/transformers/models/lfm2_vl/processing_lfm2_vl.py b/third_party/transformers/src/transformers/models/lfm2_vl/processing_lfm2_vl.py new file mode 100644 index 0000000000000000000000000000000000000000..bf654310d0d3d45b11d8c901265129472e2384ec --- /dev/null +++ b/third_party/transformers/src/transformers/models/lfm2_vl/processing_lfm2_vl.py @@ -0,0 +1,272 @@ +# 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 math + +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput, make_nested_list_of_images +from ...processing_utils import ( + ProcessingKwargs, + ProcessorMixin, + TextKwargs, + Unpack, +) +from ...tokenization_utils_base import BatchEncoding, TextInput +from ...utils import auto_docstring, logging + + +logger = logging.get_logger(__name__) + + +class Lfm2VlTextKwargs(TextKwargs, total=False): + """ + use_image_special_tokens (`bool`, *optional*, defaults to `True`): + Whether to use special image tokens (`<|image_start|>` and `<|image_end|>`) to delimit image sequences + in the text. When enabled, images are wrapped with these tokens to clearly mark image boundaries. + When disabled, only the image token itself is used without delimiters. + """ + + use_image_special_tokens: bool | None + + +class Lfm2VlProcessorKwargs(ProcessingKwargs, total=False): + text_kwargs: Lfm2VlTextKwargs + _defaults = { + "images_kwargs": { + "return_row_col_info": True, + }, + "text_kwargs": { + "use_image_special_tokens": True, + "add_special_tokens": False, + "padding": False, + "is_split_into_words": False, + }, + } + + +@auto_docstring +class Lfm2VlProcessor(ProcessorMixin): + def __init__( + self, + image_processor, + tokenizer, + chat_template: str | None = None, + **kwargs, + ): + self.image_token = getattr(tokenizer, "image_token", "") + self.image_token_id = ( + tokenizer.image_token_id + if hasattr(tokenizer, "image_token_id") + else tokenizer.convert_tokens_to_ids(self.image_token) + ) + self.image_start_token = getattr(tokenizer, "image_start_token", "<|image_start|>") + self.image_end_token = getattr(tokenizer, "image_end_token", "<|image_end|>") + self.image_thumbnail_token = getattr(tokenizer, "image_thumbnail_token", "<|img_thumbnail|>") + super().__init__(image_processor, tokenizer, chat_template=chat_template, **kwargs) + + @auto_docstring + def __call__( + self, + images: ImageInput | list[ImageInput] | list[list[ImageInput]] | None = None, + text: TextInput | list[TextInput] | None = None, + **kwargs: Unpack[Lfm2VlProcessorKwargs], + ) -> BatchEncoding: + if text is None and images is None: + raise ValueError("You must provide one of `text` or `images`.") + + if images is not None and text is None: + raise ValueError( + "You must provide `text` when `images` is provided. Minimal text consists of a single image token." + ) + + output_kwargs = self._merge_kwargs( + Lfm2VlProcessorKwargs, + tokenizer_init_kwargs=self.tokenizer.init_kwargs, + **kwargs, + ) + + 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") + + n_images_in_text = [sample.count(self.image_token) for sample in text] + if sum(n_images_in_text) > 0 and images is None: + raise ValueError(f"We detected {sum(n_images_in_text)} tokens in the text but no images were passed") + + inputs = {} + use_image_special_tokens = output_kwargs["text_kwargs"].pop("use_image_special_tokens") + + if images is not None: + images = self.image_processor.fetch_images(images) + batched_images = make_nested_list_of_images(images) + vision_inputs = self.image_processor(batched_images, **output_kwargs["images_kwargs"]) + + n_images_in_images = [len(sublist) for sublist in batched_images] + if n_images_in_images != n_images_in_text: + raise ValueError( + f"The number of images in the text {n_images_in_text} and images {n_images_in_images} should be the same." + ) + + text = self.expand_text_with_placeholders( + text, + batched_images, + image_rows=vision_inputs.pop("image_rows"), + image_cols=vision_inputs.pop("image_cols"), + image_sizes=vision_inputs.pop("image_sizes"), + use_image_special_tokens=use_image_special_tokens, + **output_kwargs["images_kwargs"], + ) + inputs.update(vision_inputs) + + return_tensors = output_kwargs["text_kwargs"].pop("return_tensors", None) + + text_inputs = self.tokenizer(text, **output_kwargs["text_kwargs"]) + inputs.update(text_inputs) + + return BatchFeature(inputs, tensor_type=return_tensors) + + def expand_text_with_placeholders( + self, + text: list[str], + images: list[list[ImageInput]], + image_rows: list[list[int]], + image_cols: list[list[int]], + image_sizes: list[list[int]], + use_image_special_tokens: bool, + **images_kwargs, + ) -> list[str]: + use_thumbnail = images_kwargs.get("use_thumbnail", self.image_processor.use_thumbnail) + image_data = iter(zip(image_rows, image_cols, image_sizes)) + + prompt_strings = [] + for sample_text, sample_images in zip(text, images): + text_parts = sample_text.split(self.image_token) + result_parts = [] + + for i, _ in enumerate(sample_images): + result_parts.append(text_parts[i]) + + rows, cols, image_size = next(image_data) + tokens_per_tile, tokens_for_image = self._get_image_num_tokens(image_size, **images_kwargs) + image_tokens = self._build_image_tokens( + rows, + cols, + tokens_per_tile, + tokens_for_image, + use_thumbnail, + use_image_special_tokens, + ) + result_parts.append(image_tokens) + + # Add remaining text after the last image + if len(sample_images) < len(text_parts): + result_parts.append(text_parts[-1]) + + prompt_strings.append("".join(result_parts)) + + return prompt_strings + + def _build_image_tokens( + self, + rows: int, + cols: int, + tokens_per_tile: int, + tokens_for_image: int, + use_thumbnail: bool, + use_image_special_tokens: bool, + ) -> str: + """Build the expanded token string for a single image.""" + parts = [] + + if use_image_special_tokens: + parts.append(self.image_start_token) + + is_multi_tile = rows > 1 or cols > 1 + if is_multi_tile: + for row in range(rows): + for col in range(cols): + if use_image_special_tokens: + parts.append(f"<|img_row_{row + 1}_col_{col + 1}|>") + parts.append(self.image_token * tokens_per_tile) + + if use_thumbnail: + if use_image_special_tokens: + parts.append(self.image_thumbnail_token) + parts.append(self.image_token * tokens_for_image) + else: + parts.append(self.image_token * tokens_for_image) + + if use_image_special_tokens: + parts.append(self.image_end_token) + + return "".join(parts) + + def _compute_tokens_per_tile(self, tile_size: int, encoder_patch_size: int, downsample_factor: int) -> int: + """Compute the number of tokens for a single tile.""" + num_patches = tile_size // encoder_patch_size + downsampled_patches = math.ceil(num_patches / downsample_factor) + return downsampled_patches * downsampled_patches + + def _compute_tokens_for_image(self, image_size: list[int], encoder_patch_size: int, downsample_factor: int) -> int: + """Compute the number of tokens for a resized image (used for single-tile or thumbnail).""" + image_height, image_width = image_size + patches_h = math.ceil((image_height // encoder_patch_size) / downsample_factor) + patches_w = math.ceil((image_width // encoder_patch_size) / downsample_factor) + return patches_h * patches_w + + def _get_image_num_tokens(self, image_size: list[int], **images_kwargs) -> tuple[int, int]: + """ + Compute token counts for image processing. + + Returns: + tuple[int, int]: (tokens_per_tile, tokens_for_image) + - tokens_per_tile: tokens for each tile in multi-tile mode + - tokens_for_image: tokens for the resized image (single-tile) or thumbnail (multi-tile) + """ + tile_size = images_kwargs.get("tile_size", self.image_processor.tile_size) + downsample_factor = images_kwargs.get("downsample_factor", self.image_processor.downsample_factor) + encoder_patch_size = images_kwargs.get("encoder_patch_size", self.image_processor.encoder_patch_size) + + tokens_per_tile = self._compute_tokens_per_tile(tile_size, encoder_patch_size, downsample_factor) + tokens_for_image = self._compute_tokens_for_image(image_size, encoder_patch_size, downsample_factor) + + return tokens_per_tile, tokens_for_image + + def batch_decode(self, *args, **kwargs): + """ + This method forwards all its arguments to LFM2Tokeniser's [`~PreTrainedTokenizer.batch_decode`]. Please + refer to the docstring of this method for more information. + """ + batched_decode_output = self.tokenizer.batch_decode(*args, **kwargs) + return batched_decode_output + + def decode(self, *args, **kwargs): + """ + This method forwards all its arguments to LFM2Tokeniser's [`~PreTrainedTokenizer.decode`]. Please refer to + the docstring of this method for more information. + """ + decode_output = self.tokenizer.decode(*args, **kwargs) + return decode_output + + @property + def model_input_names(self): + tokenizer_input_names = self.tokenizer.model_input_names + image_processor_input_names = self.image_processor.model_input_names + + # LFM2-VL has no dedicated tokenizer class and uses the Base class with default model input names + tokenizer_input_names = [name for name in tokenizer_input_names if name != "token_type_ids"] + return list(tokenizer_input_names + image_processor_input_names) + + +__all__ = ["Lfm2VlProcessor"] diff --git a/third_party/transformers/src/transformers/models/minimax_m2/__init__.py b/third_party/transformers/src/transformers/models/minimax_m2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..87f3fe717035c8f6541717add547a5b7644107d7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/minimax_m2/__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_minimax_m2 import * + from .modeling_minimax_m2 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/minimax_m2/configuration_minimax_m2.py b/third_party/transformers/src/transformers/models/minimax_m2/configuration_minimax_m2.py new file mode 100644 index 0000000000000000000000000000000000000000..75acb8b755d785c9b00e4dd344ba086459e7121d --- /dev/null +++ b/third_party/transformers/src/transformers/models/minimax_m2/configuration_minimax_m2.py @@ -0,0 +1,94 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/minimax_m2/modular_minimax_m2.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_minimax_m2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the MiniMax AI Team and 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="MiniMaxAI/MiniMax-Text-01-hf") +@strict +class MiniMaxM2Config(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import MiniMaxM2Model, MiniMaxM2Config + + >>> # Initializing a MiniMaxM2 style configuration + >>> configuration = MiniMaxM2Config() + + >>> # Initializing a model from the MiniMaxM2 style configuration + >>> model = MiniMaxM2Model(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "minimax_m2" + keys_to_ignore_at_inference = ["past_key_values"] + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise_gather_output", + "layers.*.self_attn.k_proj": "colwise_gather_output", + "layers.*.self_attn.v_proj": "colwise_gather_output", + "layers.*.self_attn.o_proj": "rowwise_split_input", + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + } + 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_experts": "num_local_experts", + } + default_theta = 5000000.0 + + vocab_size: int = 200064 + hidden_size: int = 3072 + intermediate_size: int = 1536 + num_hidden_layers: int = 62 + num_attention_heads: int = 48 + num_key_value_heads: int = 8 + head_dim: int = 128 + hidden_act: str = "silu" + max_position_embeddings: int = 196608 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-06 + use_cache: bool = True + pad_token_id: int | None = None + bos_token_id: int | None = 200034 + eos_token_id: int | list[int] | None = 200020 + tie_word_embeddings: bool = False + attention_dropout: float | int = 0.0 + num_experts_per_tok: int = 8 + num_local_experts: int = 256 + output_router_logits: bool = False + router_aux_loss_coef: float = 0.001 + router_jitter_noise: float = 0.0 + rope_parameters: RopeParameters | dict | None = None + + +__all__ = ["MiniMaxM2Config"] diff --git a/third_party/transformers/src/transformers/models/minimax_m2/modeling_minimax_m2.py b/third_party/transformers/src/transformers/models/minimax_m2/modeling_minimax_m2.py new file mode 100644 index 0000000000000000000000000000000000000000..d19274262810cf8699b38b9fdd0fe36996d4aebf --- /dev/null +++ b/third_party/transformers/src/transformers/models/minimax_m2/modeling_minimax_m2.py @@ -0,0 +1,691 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/minimax_m2/modular_minimax_m2.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_minimax_m2.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the MiniMax AI Team and 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.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_experts_implementation, use_kernel_forward_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, can_return_tuple +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from .configuration_minimax_m2 import MiniMaxM2Config + + +class MiniMaxM2TopKRouter(nn.Module): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_local_experts + self.hidden_dim = config.hidden_size + self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim)) + + def forward(self, hidden_states, e_score_correction_bias): + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + router_logits = F.linear(hidden_states.to(self.weight.dtype), self.weight) # (seq_len, num_experts) + # Main difference to other Moe, using Sigmoid activation instead of Softmax + routing_weights = nn.functional.sigmoid(router_logits.float()) + scores_for_choice = routing_weights + e_score_correction_bias + _, top_k_index = torch.topk(scores_for_choice, self.top_k, dim=-1, sorted=False) + top_k_weights = routing_weights.gather(1, top_k_index) + top_k_weights /= top_k_weights.sum(dim=-1, keepdim=True) + router_scores = top_k_weights + return router_logits, router_scores, top_k_index + + +@use_experts_implementation +class MiniMaxM2Experts(nn.Module): + """Collection of expert weights stored as 3D tensors.""" + + def __init__(self, config: MiniMaxM2Config): + super().__init__() + self.num_experts = config.num_local_experts + self.hidden_dim = config.hidden_size + self.intermediate_dim = config.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 MiniMaxM2SparseMoeBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.jitter_noise = config.router_jitter_noise + self.gate = MiniMaxM2TopKRouter(config) + self.experts = MiniMaxM2Experts(config) + self.register_buffer("e_score_correction_bias", torch.zeros(config.num_local_experts)) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, sequence_length, hidden_dim = hidden_states.shape + if self.training and self.jitter_noise > 0: + hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + _, top_k_weights, top_k_index = self.gate(hidden_states, self.e_score_correction_bias) + hidden_states = self.experts(hidden_states, top_k_index, top_k_weights) + hidden_states = hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return hidden_states + + +@use_kernel_forward_from_hub("RMSNorm") +class MiniMaxM2RMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + MiniMaxM2RMSNorm 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 MiniMaxM2RotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: MiniMaxM2Config, 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: MiniMaxM2Config | 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 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(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 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_kernelized_func(apply_rotary_pos_emb) +class MiniMaxM2Attention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: MiniMaxM2Config, 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=False) + self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) + self.q_norm = MiniMaxM2RMSNorm(config.num_attention_heads * self.head_dim, config.rms_norm_eps) + self.k_norm = MiniMaxM2RMSNorm(config.num_key_value_heads * self.head_dim, 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 | None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_norm(self.q_proj(hidden_states)) + key_states = self.k_norm(self.k_proj(hidden_states)) + value_states = self.v_proj(hidden_states) + + query_states = query_states.view(hidden_shape).transpose(1, 2) + key_states = key_states.view(hidden_shape).transpose(1, 2) + value_states = value_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 MiniMaxM2DecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: MiniMaxM2Config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = MiniMaxM2Attention(config, layer_idx) + + self.mlp = MiniMaxM2SparseMoeBlock(config) + self.input_layernorm = MiniMaxM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = MiniMaxM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + 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, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + 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, + **kwargs, + ) + hidden_states = residual + hidden_states + 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 MiniMaxM2PreTrainedModel(PreTrainedModel): + config: MiniMaxM2Config + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["MiniMaxM2DecoderLayer"] + _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 = { + "router_logits": OutputRecorder(MiniMaxM2TopKRouter, index=0), + "hidden_states": MiniMaxM2DecoderLayer, + "attentions": MiniMaxM2Attention, + } + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + std = self.config.initializer_range + if isinstance(module, MiniMaxM2Experts): + init.normal_(module.gate_up_proj, mean=0.0, std=std) + init.normal_(module.down_proj, mean=0.0, std=std) + elif isinstance(module, MiniMaxM2TopKRouter): + init.normal_(module.weight, mean=0.0, std=std) + elif isinstance(module, MiniMaxM2SparseMoeBlock): + init.zeros_(module.e_score_correction_bias) + + +@auto_docstring +class MiniMaxM2Model(MiniMaxM2PreTrainedModel): + def __init__(self, config: MiniMaxM2Config): + 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( + [MiniMaxM2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = MiniMaxM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = MiniMaxM2RotaryEmbedding(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], + ) -> 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) + + # No sliding window opposed to mixtral + 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_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **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 MiniMaxM2ForCausalLM(MiniMaxM2PreTrainedModel, 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 = MiniMaxM2Model(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 + + # 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, + output_router_logits: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> 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, MiniMaxM2ForCausalLM + + >>> model = MiniMaxM2ForCausalLM.from_pretrained("mistralai/MiniMaxM2-8x7B-v0.1") + >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/MiniMaxM2-8x7B-v0.1") + + >>> 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: MoeModelOutputWithPast = 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, + output_router_logits=output_router_logits, + **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, labels, self.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__ = ["MiniMaxM2ForCausalLM", "MiniMaxM2Model", "MiniMaxM2PreTrainedModel"] diff --git a/third_party/transformers/src/transformers/models/minimax_m2/modular_minimax_m2.py b/third_party/transformers/src/transformers/models/minimax_m2/modular_minimax_m2.py new file mode 100644 index 0000000000000000000000000000000000000000..a9938a555c62ab276833951a429562c4e570a9b2 --- /dev/null +++ b/third_party/transformers/src/transformers/models/minimax_m2/modular_minimax_m2.py @@ -0,0 +1,246 @@ +# Copyright 2025 the MiniMax AI Team and 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 +import torch.nn.functional as F +from huggingface_hub.dataclasses import strict +from torch import nn + +from ... import initialization as init +from ...cache_utils import Cache, DynamicCache +from ...configuration_utils import PreTrainedConfig +from ...masking_utils import create_causal_mask +from ...modeling_outputs import MoeModelOutputWithPast +from ...modeling_rope_utils import RopeParameters +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, auto_docstring +from ...utils.generic import merge_with_config_defaults +from ...utils.output_capturing import capture_outputs +from ..flex_olmo.modeling_flex_olmo import FlexOlmoAttention +from ..glm4_moe.modeling_glm4_moe import ( + Glm4MoeRotaryEmbedding, + apply_rotary_pos_emb, # noqa: F401 +) +from ..mixtral.modeling_mixtral import ( + MixtralExperts, + MixtralForCausalLM, + MixtralModel, + MixtralPreTrainedModel, + MixtralRMSNorm, + MixtralSparseMoeBlock, + MixtralTopKRouter, +) + + +@auto_docstring(checkpoint="MiniMaxAI/MiniMax-Text-01-hf") +@strict +class MiniMaxM2Config(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import MiniMaxM2Model, MiniMaxM2Config + + >>> # Initializing a MiniMaxM2 style configuration + >>> configuration = MiniMaxM2Config() + + >>> # Initializing a model from the MiniMaxM2 style configuration + >>> model = MiniMaxM2Model(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "minimax_m2" + keys_to_ignore_at_inference = ["past_key_values"] + base_model_tp_plan = { + "layers.*.self_attn.q_proj": "colwise_gather_output", + "layers.*.self_attn.k_proj": "colwise_gather_output", + "layers.*.self_attn.v_proj": "colwise_gather_output", + "layers.*.self_attn.o_proj": "rowwise_split_input", + "layers.*.mlp.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + } + 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_experts": "num_local_experts", + } + default_theta = 5000000.0 + + vocab_size: int = 200064 + hidden_size: int = 3072 + intermediate_size: int = 1536 + num_hidden_layers: int = 62 + num_attention_heads: int = 48 + num_key_value_heads: int = 8 + head_dim: int = 128 + hidden_act: str = "silu" + max_position_embeddings: int = 196608 + initializer_range: float = 0.02 + rms_norm_eps: float = 1e-06 + use_cache: bool = True + pad_token_id: int | None = None + bos_token_id: int | None = 200034 + eos_token_id: int | list[int] | None = 200020 + tie_word_embeddings: bool = False + attention_dropout: float | int = 0.0 + num_experts_per_tok: int = 8 + num_local_experts: int = 256 + output_router_logits: bool = False + router_aux_loss_coef: float = 0.001 + router_jitter_noise: float = 0.0 + rope_parameters: RopeParameters | dict | None = None + + +class MiniMaxM2TopKRouter(MixtralTopKRouter): + def forward(self, hidden_states, e_score_correction_bias): + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + router_logits = F.linear(hidden_states.to(self.weight.dtype), self.weight) # (seq_len, num_experts) + # Main difference to other Moe, using Sigmoid activation instead of Softmax + routing_weights = nn.functional.sigmoid(router_logits.float()) + scores_for_choice = routing_weights + e_score_correction_bias + _, top_k_index = torch.topk(scores_for_choice, self.top_k, dim=-1, sorted=False) + top_k_weights = routing_weights.gather(1, top_k_index) + top_k_weights /= top_k_weights.sum(dim=-1, keepdim=True) + router_scores = top_k_weights + return router_logits, router_scores, top_k_index + + +class MiniMaxM2Experts(MixtralExperts): + pass + + +class MiniMaxM2SparseMoeBlock(MixtralSparseMoeBlock): + def __init__(self, config): + super().__init__() + self.register_buffer("e_score_correction_bias", torch.zeros(config.num_local_experts)) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, sequence_length, hidden_dim = hidden_states.shape + if self.training and self.jitter_noise > 0: + hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + _, top_k_weights, top_k_index = self.gate(hidden_states, self.e_score_correction_bias) + hidden_states = self.experts(hidden_states, top_k_index, top_k_weights) + hidden_states = hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return hidden_states + + +class MiniMaxM2RMSNorm(MixtralRMSNorm): + pass + + +class MiniMaxM2RotaryEmbedding(Glm4MoeRotaryEmbedding): + pass + + +class MiniMaxM2Attention(FlexOlmoAttention): + def __init__(self, config: MiniMaxM2Config, layer_idx: int): + super().__init__(config, layer_idx) + self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False) + self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) + + +class MiniMaxM2PreTrainedModel(MixtralPreTrainedModel): + @torch.no_grad() + def _init_weights(self, module): + PreTrainedModel._init_weights(self, module) + std = self.config.initializer_range + if isinstance(module, MiniMaxM2Experts): + init.normal_(module.gate_up_proj, mean=0.0, std=std) + init.normal_(module.down_proj, mean=0.0, std=std) + elif isinstance(module, MiniMaxM2TopKRouter): + init.normal_(module.weight, mean=0.0, std=std) + elif isinstance(module, MiniMaxM2SparseMoeBlock): + init.zeros_(module.e_score_correction_bias) + + +class MiniMaxM2Model(MixtralModel): + @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) + + # No sliding window opposed to mixtral + 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_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **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, + ) + + +class MiniMaxM2ForCausalLM(MixtralForCausalLM): + pass + + +__all__ = [ + "MiniMaxM2Config", + "MiniMaxM2ForCausalLM", + "MiniMaxM2Model", # noqa: F822 + "MiniMaxM2PreTrainedModel", # noqa: F822 +] diff --git a/third_party/transformers/src/transformers/models/mixtral/__init__.py b/third_party/transformers/src/transformers/models/mixtral/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e4ca36bacbee79261b8fe19e1df5cd2fb91b7344 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mixtral/__init__.py @@ -0,0 +1,27 @@ +# Copyright 2023 Mixtral 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. +from typing import TYPE_CHECKING + +from ...utils import _LazyModule +from ...utils.import_utils import define_import_structure + + +if TYPE_CHECKING: + from .configuration_mixtral import * + from .modeling_mixtral 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/mixtral/configuration_mixtral.py b/third_party/transformers/src/transformers/models/mixtral/configuration_mixtral.py new file mode 100644 index 0000000000000000000000000000000000000000..240f2441103153f68fed7b631aee132b11932126 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mixtral/configuration_mixtral.py @@ -0,0 +1,93 @@ +# Copyright 2023 Mixtral 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. +"""Mixtral 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="mistralai/Mixtral-8x7B-v0.1") +@strict +class MixtralConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import MixtralModel, MixtralConfig + + >>> # Initializing a Mixtral 7B style configuration + >>> configuration = MixtralConfig() + + >>> # Initializing a model from the Mixtral 7B style configuration + >>> model = MixtralModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "mixtral" + keys_to_ignore_at_inference = ["past_key_values"] + default_theta = 1000000.0 + 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.experts.gate_up_proj": "packed_colwise", + "layers.*.mlp.experts.down_proj": "rowwise", + "layers.*.mlp.experts": "moe_tp_experts", + } + 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_experts": "num_local_experts"} + + vocab_size: int = 32000 + hidden_size: int = 4096 + intermediate_size: int = 14336 + num_hidden_layers: int = 32 + num_attention_heads: int = 32 + num_key_value_heads: int = 8 + head_dim: int | None = None + hidden_act: str = "silu" + max_position_embeddings: int = 4096 * 32 + 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 = 1 + eos_token_id: int | list[int] | None = 2 + tie_word_embeddings: bool = False + sliding_window: int | None = None + attention_dropout: float | int = 0.0 + num_experts_per_tok: int = 2 + num_local_experts: int = 8 + output_router_logits: bool = False + router_aux_loss_coef: float = 0.001 + router_jitter_noise: float = 0.0 + rope_parameters: RopeParameters | dict | 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__ = ["MixtralConfig"] diff --git a/third_party/transformers/src/transformers/models/mixtral/convert_mixtral_weights_to_hf.py b/third_party/transformers/src/transformers/models/mixtral/convert_mixtral_weights_to_hf.py new file mode 100644 index 0000000000000000000000000000000000000000..0ae7cdd5475ca5bed9e762afda47cc333c3e2c9a --- /dev/null +++ b/third_party/transformers/src/transformers/models/mixtral/convert_mixtral_weights_to_hf.py @@ -0,0 +1,243 @@ +# Copyright 2023 Mistral 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. +import argparse +import json +import os + +import torch + +from transformers import ( + MixtralConfig, + MixtralForCausalLM, +) + + +""" +Sample usage: + +``` +python src/transformers/models/mixtral/convert_mixtral_weights_to_hf.py \ + --input_dir /path/to/downloaded/mixtral/weights --model_size 7B --output_dir /output/path +``` + +Thereafter, models can be loaded via: + +```py +from transformers import MixtralForCausalLM + +model = MixtralForCausalLM.from_pretrained("/output/path") +``` + +Important note: you need to be able to host the whole model in RAM to execute this script (even if the biggest versions +come in several checkpoints they each contain a part of each weight of the model, so we need to load them all in RAM). +""" + + +def compute_intermediate_size(n, ffn_dim_multiplier=1, multiple_of=256): + return multiple_of * ((int(ffn_dim_multiplier * int(8 * n / 3)) + multiple_of - 1) // multiple_of) + + +def read_json(path): + with open(path, "r") as f: + return json.load(f) + + +def write_json(text, path): + with open(path, "w") as f: + json.dump(text, f) + + +def write_model(model_path, input_base_path, model_size): + os.makedirs(model_path, exist_ok=True) + + params = read_json(os.path.join(input_base_path, "params.json")) + num_shards = 1 + + # For some reason this is a string in the params.json + sliding_window = int(params["sliding_window"]) if "sliding_window" in params else None + n_layers = params["num_hidden_layers"] + n_heads = params["num_attention_heads"] + n_heads_per_shard = n_heads // num_shards + dim = params["hidden_size"] + dims_per_head = dim // n_heads + base = params.get("rope_theta", 10000.0) + max_position_embeddings = 4096 * 8 + num_local_experts = params["num_local_experts"] + ffn_dim = params["intermediate_size"] + + vocab_size = params["vocab_size"] + + if "num_key_value_heads" in params: + num_key_value_heads = params["num_key_value_heads"] # for GQA / MQA + num_local_key_value_heads = num_key_value_heads // num_shards + key_value_dim = dims_per_head * num_local_key_value_heads + else: # compatibility with other checkpoints + num_key_value_heads = n_heads + num_local_key_value_heads = n_heads_per_shard + key_value_dim = dim + + # permute for sliced rotary + def permute(w, n_heads=n_heads, dim1=dim, dim2=dim): + return w.view(n_heads, dim1 // n_heads // 2, 2, dim2).transpose(1, 2).reshape(dim1, dim2) + + print(f"Fetching all parameters from the checkpoint at {input_base_path}.") + # Load weights + loaded = [ + torch.load(os.path.join(input_base_path, f"consolidated.{i:02d}.pt"), map_location="cpu", weights_only=True) + for i in range(8) + ] + + merged_state_dict = {} + for state_dict in loaded: + merged_state_dict.update(state_dict) + + state_dict = {} + + for layer_i in range(n_layers): + # Sharded + # Note that attention.w{q,k,v,o}, feed_fordward.w[1,2,3], attention_norm.weight and ffn_norm.weight share + # the same storage object, saving attention_norm and ffn_norm will save other weights too, which is + # redundant as other weights will be stitched from multiple shards. To avoid that, they are cloned. + + state_dict.update( + { + f"model.layers.{layer_i}.input_layernorm.weight": merged_state_dict[ + f"layers.{layer_i}.attention_norm.weight" + ].clone(), + f"model.layers.{layer_i}.post_attention_layernorm.weight": merged_state_dict[ + f"layers.{layer_i}.ffn_norm.weight" + ].clone(), + } + ) + + state_dict[f"model.layers.{layer_i}.self_attn.q_proj.weight"] = permute( + merged_state_dict[f"layers.{layer_i}.attention.wq.weight"] + .view(n_heads_per_shard, dims_per_head, dim) + .reshape(dim, dim) + ) + state_dict[f"model.layers.{layer_i}.self_attn.k_proj.weight"] = permute( + merged_state_dict[f"layers.{layer_i}.attention.wk.weight"] + .view(num_local_key_value_heads, dims_per_head, dim) + .reshape(key_value_dim, dim), + num_key_value_heads, + key_value_dim, + dim, + ) + state_dict[f"model.layers.{layer_i}.self_attn.v_proj.weight"] = ( + merged_state_dict[f"layers.{layer_i}.attention.wv.weight"] + .view(num_local_key_value_heads, dims_per_head, dim) + .reshape(key_value_dim, dim) + ) + + state_dict[f"model.layers.{layer_i}.self_attn.o_proj.weight"] = merged_state_dict[ + f"layers.{layer_i}.attention.wo.weight" + ] + + w1 = merged_state_dict[f"layers.{layer_i}.block_sparse_moe.w1"] + w2 = merged_state_dict[f"layers.{layer_i}.block_sparse_moe.w2"] + w3 = merged_state_dict[f"layers.{layer_i}.block_sparse_moe.w3"] + + experts_w1 = [ + w1[ffn_dim * expert_idx : ffn_dim * (expert_idx + 1), :].clone(memory_format=torch.contiguous_format) + for expert_idx in range(num_local_experts) + ] + + for idx, expert_block in enumerate(experts_w1): + expert_key = f"model.layers.{layer_i}.block_sparse_moe.experts.{idx}.w1" + state_dict[expert_key + ".weight"] = expert_block.clone() + + experts_w2 = [ + w2[ffn_dim * expert_idx : ffn_dim * (expert_idx + 1), :].clone(memory_format=torch.contiguous_format) + for expert_idx in range(num_local_experts) + ] + + for idx, expert_block in enumerate(experts_w2): + expert_key = f"model.layers.{layer_i}.block_sparse_moe.experts.{idx}.w2" + state_dict[expert_key + ".weight"] = expert_block.T.clone(memory_format=torch.contiguous_format) + + experts_w3 = [ + w3[ffn_dim * expert_idx : ffn_dim * (expert_idx + 1), :].clone(memory_format=torch.contiguous_format) + for expert_idx in range(num_local_experts) + ] + + for idx, expert_block in enumerate(experts_w3): + expert_key = f"model.layers.{layer_i}.block_sparse_moe.experts.{idx}.w3" + state_dict[expert_key + ".weight"] = expert_block.clone() + + state_dict[f"model.layers.{layer_i}.block_sparse_moe.gate.weight"] = merged_state_dict[ + f"layers.{layer_i}.block_sparse_moe.gate.weight" + ] + + state_dict.update( + { + "model.norm.weight": merged_state_dict["norm.weight"], + "model.embed_tokens.weight": merged_state_dict["tok_embeddings.weight"], + "lm_head.weight": merged_state_dict["output.weight"], + } + ) + + config = MixtralConfig( + hidden_size=dim, + intermediate_size=ffn_dim, + num_attention_heads=params["num_attention_heads"], + num_hidden_layers=params["num_hidden_layers"], + rms_norm_eps=params["rms_norm_eps"], + num_key_value_heads=num_key_value_heads, + vocab_size=vocab_size, + rope_theta=base, + max_position_embeddings=max_position_embeddings, + sliding_window=sliding_window, + num_local_experts=num_local_experts, + ) + + print("Loading the checkpoint in a Mixtral model.") + with torch.device("meta"): + model = MixtralForCausalLM(config) + # Avoid saving this as part of the config. + del model.config._name_or_path + model.config.dtype = torch.float16 + print("Saving in the Transformers format.") + + model.load_state_dict(state_dict, strict=True, assign=True) + + for n, p in model.named_parameters(): + assert p.device.type != "meta", f"{n} has not been loaded!" + + model.save_pretrained(model_path) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_dir", + help="Location of Mixtral weights, which contains tokenizer.model and model folders", + required=True, + ) + parser.add_argument( + "--model_size", + choices=["7B"], + help="'f' models correspond to the finetuned versions, and are specific to the Mixtral official release. For more details on Mixtral, check out the original repo: https://huggingface.co/mistral-ai", + default="7B", + ) + parser.add_argument("--output_dir", help="Location to write HF model", required=True) + args = parser.parse_args() + write_model( + model_path=args.output_dir, + input_base_path=args.input_dir, + model_size=args.model_size, + ) + + +if __name__ == "__main__": + main() diff --git a/third_party/transformers/src/transformers/models/mixtral/modeling_mixtral.py b/third_party/transformers/src/transformers/models/mixtral/modeling_mixtral.py new file mode 100644 index 0000000000000000000000000000000000000000..3c75687c4c49b87f50bf1f88244bcb3fd99d90f7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mixtral/modeling_mixtral.py @@ -0,0 +1,702 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/mixtral/modular_mixtral.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_mixtral.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2023 Mistral AI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +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_experts_implementation, + use_kernel_forward_from_hub, + use_kernel_func_from_hub, + use_kernelized_func, +) +from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import ( + GenericForQuestionAnswering, + GenericForSequenceClassification, + GenericForTokenClassification, + 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, can_return_tuple +from ...utils.generic import maybe_autocast, merge_with_config_defaults +from ...utils.output_capturing import OutputRecorder, capture_outputs +from .configuration_mixtral import MixtralConfig + + +@use_experts_implementation +class MixtralExperts(nn.Module): + """Collection of expert weights stored as 3D tensors.""" + + def __init__(self, config: MixtralConfig): + super().__init__() + self.num_experts = config.num_local_experts + self.hidden_dim = config.hidden_size + self.intermediate_dim = config.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 MixtralTopKRouter(nn.Module): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_local_experts + self.hidden_dim = config.hidden_size + self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim)) + + def forward(self, hidden_states): + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + router_logits = F.linear(hidden_states, self.weight) # (seq_len, num_experts) + router_logits = torch.nn.functional.softmax(router_logits.float(), dim=-1) + router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1) # (seq_len, top_k) + router_top_value /= router_top_value.sum(dim=-1, keepdim=True) + router_scores = router_top_value + return router_logits, router_scores, router_indices + + +class MixtralSparseMoeBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.jitter_noise = config.router_jitter_noise + self.gate = MixtralTopKRouter(config) + self.experts = MixtralExperts(config) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, sequence_length, hidden_dim = hidden_states.shape + if self.training and self.jitter_noise > 0: + hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + _, top_k_weights, top_k_index = self.gate(hidden_states) + hidden_states = self.experts(hidden_states, top_k_index, top_k_weights) + hidden_states = hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return hidden_states + + +@use_kernel_forward_from_hub("RMSNorm") +class MixtralRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + MixtralRMSNorm 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 MixtralRotaryEmbedding(nn.Module): + inv_freq: torch.Tensor # fix linting for `register_buffer` + + def __init__(self, config: MixtralConfig, 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: MixtralConfig | 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 MixtralAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: MixtralConfig, layer_idx: int): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.head_dim = getattr(config, "head_dim", None) or 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=False) + self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False) + 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], + 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) + 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, + sliding_window=getattr(self.config, "sliding_window", None), # main diff with Llama + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class MixtralDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: MixtralConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = MixtralAttention(config, layer_idx) + + self.mlp = MixtralSparseMoeBlock(config) + self.input_layernorm = MixtralRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = MixtralRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + 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, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + 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, + **kwargs, + ) + hidden_states = residual + hidden_states + 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 MixtralPreTrainedModel(PreTrainedModel): + config: MixtralConfig + base_model_prefix = "model" + supports_gradient_checkpointing = True + _no_split_modules = ["MixtralDecoderLayer"] + _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 = { + "router_logits": OutputRecorder(MixtralTopKRouter, index=0), + "hidden_states": MixtralDecoderLayer, + "attentions": MixtralAttention, + } + + @torch.no_grad() + def _init_weights(self, module): + super()._init_weights(module) + std = self.config.initializer_range + if isinstance(module, MixtralExperts): + init.normal_(module.gate_up_proj, mean=0.0, std=std) + init.normal_(module.down_proj, mean=0.0, std=std) + elif isinstance(module, MixtralTopKRouter): + init.normal_(module.weight, mean=0.0, std=std) + + +@auto_docstring +class MixtralModel(MixtralPreTrainedModel): + def __init__(self, config: MixtralConfig): + 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( + [MixtralDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = MixtralRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = MixtralRotaryEmbedding(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], + ) -> 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) + + mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask + causal_mask = mask_function( + 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_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **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 MixtralForCausalLM(MixtralPreTrainedModel, 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 = MixtralModel(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 + + # 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, + output_router_logits: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> 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, MixtralForCausalLM + + >>> model = MixtralForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-v0.1") + >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-v0.1") + + >>> 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: MoeModelOutputWithPast = 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, + output_router_logits=output_router_logits, + **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, labels, self.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, + ) + + +class MixtralForSequenceClassification(GenericForSequenceClassification, MixtralPreTrainedModel): + pass + + +class MixtralForTokenClassification(GenericForTokenClassification, MixtralPreTrainedModel): + pass + + +class MixtralForQuestionAnswering(GenericForQuestionAnswering, MixtralPreTrainedModel): + pass + + +__all__ = [ + "MixtralForCausalLM", + "MixtralForQuestionAnswering", + "MixtralModel", + "MixtralPreTrainedModel", + "MixtralForSequenceClassification", + "MixtralForTokenClassification", +] diff --git a/third_party/transformers/src/transformers/models/mixtral/modular_mixtral.py b/third_party/transformers/src/transformers/models/mixtral/modular_mixtral.py new file mode 100644 index 0000000000000000000000000000000000000000..2ec3d29a999b97a58576c297fbe6630f07337a09 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mixtral/modular_mixtral.py @@ -0,0 +1,448 @@ +# Copyright 2023 Mistral AI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""PyTorch Mixtral model.""" + +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 import use_experts_implementation +from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast +from ...modeling_utils import PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, logging +from ...utils.output_capturing import OutputRecorder +from ..mistral.modeling_mistral import ( + MistralAttention, + MistralForCausalLM, + MistralForQuestionAnswering, + MistralForSequenceClassification, + MistralForTokenClassification, + MistralModel, + MistralPreTrainedModel, + MistralRMSNorm, + MistralRotaryEmbedding, +) +from .configuration_mixtral import MixtralConfig + + +logger = logging.get_logger(__name__) + + +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 + + +@use_experts_implementation +class MixtralExperts(nn.Module): + """Collection of expert weights stored as 3D tensors.""" + + def __init__(self, config: MixtralConfig): + super().__init__() + self.num_experts = config.num_local_experts + self.hidden_dim = config.hidden_size + self.intermediate_dim = config.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 MixtralTopKRouter(nn.Module): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.num_experts = config.num_local_experts + self.hidden_dim = config.hidden_size + self.weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim)) + + def forward(self, hidden_states): + hidden_states = hidden_states.reshape(-1, self.hidden_dim) + router_logits = F.linear(hidden_states, self.weight) # (seq_len, num_experts) + router_logits = torch.nn.functional.softmax(router_logits.float(), dim=-1) + router_top_value, router_indices = torch.topk(router_logits, self.top_k, dim=-1) # (seq_len, top_k) + router_top_value /= router_top_value.sum(dim=-1, keepdim=True) + router_scores = router_top_value + return router_logits, router_scores, router_indices + + +class MixtralSparseMoeBlock(nn.Module): + def __init__(self, config): + super().__init__() + self.top_k = config.num_experts_per_tok + self.jitter_noise = config.router_jitter_noise + self.gate = MixtralTopKRouter(config) + self.experts = MixtralExperts(config) + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, sequence_length, hidden_dim = hidden_states.shape + if self.training and self.jitter_noise > 0: + hidden_states *= torch.empty_like(hidden_states).uniform_(1.0 - self.jitter_noise, 1.0 + self.jitter_noise) + hidden_states = hidden_states.view(-1, hidden_states.shape[-1]) + _, top_k_weights, top_k_index = self.gate(hidden_states) + hidden_states = self.experts(hidden_states, top_k_index, top_k_weights) + hidden_states = hidden_states.reshape(batch_size, sequence_length, hidden_dim) + return hidden_states + + +class MixtralRMSNorm(MistralRMSNorm): + pass + + +class MixtralRotaryEmbedding(MistralRotaryEmbedding): + pass + + +class MixtralAttention(MistralAttention): + pass + + +class MixtralDecoderLayer(GradientCheckpointingLayer): + def __init__(self, config: MixtralConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + + self.self_attn = MixtralAttention(config, layer_idx) + + self.mlp = MixtralSparseMoeBlock(config) + self.input_layernorm = MixtralRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = MixtralRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + 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, + **kwargs: Unpack[TransformersKwargs], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + 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, + **kwargs, + ) + hidden_states = residual + hidden_states + 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 + + +class MixtralPreTrainedModel(MistralPreTrainedModel): + _can_record_outputs = { + "router_logits": OutputRecorder(MixtralTopKRouter, index=0), + "hidden_states": MixtralDecoderLayer, + "attentions": MixtralAttention, + } + + @torch.no_grad() + def _init_weights(self, module): + PreTrainedModel._init_weights(self, module) + std = self.config.initializer_range + if isinstance(module, MixtralExperts): + init.normal_(module.gate_up_proj, mean=0.0, std=std) + init.normal_(module.down_proj, mean=0.0, std=std) + elif isinstance(module, MixtralTopKRouter): + init.normal_(module.weight, mean=0.0, std=std) + + +class MixtralModel(MistralModel): + 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) + + mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask + causal_mask = mask_function( + 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_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + position_embeddings=position_embeddings, + **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, + ) + + +class MixtralForCausalLM(MistralForCausalLM): + _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"} + + def __init__(self, config): + super().__init__(config) + self.model = MixtralModel(config) + 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 + + 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, + output_router_logits: bool | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> 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, MixtralForCausalLM + + >>> model = MixtralForCausalLM.from_pretrained("mistralai/Mixtral-8x7B-v0.1") + >>> tokenizer = AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-v0.1") + + >>> 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: MoeModelOutputWithPast = 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, + output_router_logits=output_router_logits, + **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, labels, self.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, + ) + + +class MixtralForSequenceClassification(MistralForSequenceClassification): + pass + + +class MixtralForTokenClassification(MistralForTokenClassification): + pass + + +class MixtralForQuestionAnswering(MistralForQuestionAnswering): + pass + + +__all__ = [ + "MixtralForCausalLM", + "MixtralForQuestionAnswering", + "MixtralModel", + "MixtralPreTrainedModel", + "MixtralForSequenceClassification", + "MixtralForTokenClassification", +] diff --git a/third_party/transformers/src/transformers/models/mra/__init__.py b/third_party/transformers/src/transformers/models/mra/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2963ad0c97ba8c0c5bdc32ff9adca79ffc56cde3 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mra/__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_mra import * + from .modeling_mra 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/mra/configuration_mra.py b/third_party/transformers/src/transformers/models/mra/configuration_mra.py new file mode 100644 index 0000000000000000000000000000000000000000..a7e2ff8ae5d5eb7b76f12985c2422e56fb289fa4 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mra/configuration_mra.py @@ -0,0 +1,76 @@ +# 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. +"""MRA model configuration""" + +from huggingface_hub.dataclasses import strict + +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring + + +@auto_docstring(checkpoint="uw-madison/mra-base-512-4") +@strict +class MraConfig(PreTrainedConfig): + r""" + block_per_row (`int`, *optional*, defaults to 4): + Used to set the budget for the high resolution scale. + approx_mode (`str`, *optional*, defaults to `"full"`): + Controls whether both low and high resolution approximations are used. Set to `"full"` for both low and + high resolution and `"sparse"` for only low resolution. + initial_prior_first_n_blocks (`int`, *optional*, defaults to 0): + The initial number of blocks for which high resolution is used. + initial_prior_diagonal_n_blocks (`int`, *optional*, defaults to 0): + The number of diagonal blocks for which high resolution is used. + + Example: + + ```python + >>> from transformers import MraConfig, MraModel + + >>> # Initializing a Mra uw-madison/mra-base-512-4 style configuration + >>> configuration = MraConfig() + + >>> # Initializing a model (with random weights) from the uw-madison/mra-base-512-4 style configuration + >>> model = MraModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "mra" + + 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 = 512 + type_vocab_size: int = 1 + initializer_range: float = 0.02 + layer_norm_eps: float = 1e-5 + block_per_row: int = 4 + approx_mode: str = "full" + initial_prior_first_n_blocks: int = 0 + initial_prior_diagonal_n_blocks: int = 0 + pad_token_id: int | None = 1 + bos_token_id: int | None = 0 + eos_token_id: int | list[int] | None = 2 + add_cross_attention: bool = False + tie_word_embeddings: bool = True + + +__all__ = ["MraConfig"] diff --git a/third_party/transformers/src/transformers/models/mra/convert_mra_pytorch_to_pytorch.py b/third_party/transformers/src/transformers/models/mra/convert_mra_pytorch_to_pytorch.py new file mode 100644 index 0000000000000000000000000000000000000000..c5873d7ad2ace3a1ce7bc7abc27682fbd6fe2bf8 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mra/convert_mra_pytorch_to_pytorch.py @@ -0,0 +1,109 @@ +# 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 MRA checkpoints from the original repository. URL: https://github.com/mlpen/mra-attention""" + +import argparse + +import torch + +from transformers import MraConfig, MraForMaskedLM + + +def rename_key(orig_key): + if "model" in orig_key: + orig_key = orig_key.replace("model.", "") + if "norm1" in orig_key: + orig_key = orig_key.replace("norm1", "attention.output.LayerNorm") + if "norm2" in orig_key: + orig_key = orig_key.replace("norm2", "output.LayerNorm") + if "norm" in orig_key: + orig_key = orig_key.replace("norm", "LayerNorm") + if "transformer" in orig_key: + layer_num = orig_key.split(".")[0].split("_")[-1] + orig_key = orig_key.replace(f"transformer_{layer_num}", f"encoder.layer.{layer_num}") + if "mha.attn" in orig_key: + orig_key = orig_key.replace("mha.attn", "attention.self") + if "mha" in orig_key: + orig_key = orig_key.replace("mha", "attention") + if "W_q" in orig_key: + orig_key = orig_key.replace("W_q", "self.query") + if "W_k" in orig_key: + orig_key = orig_key.replace("W_k", "self.key") + if "W_v" in orig_key: + orig_key = orig_key.replace("W_v", "self.value") + if "ff.0" in orig_key: + orig_key = orig_key.replace("ff.0", "intermediate.dense") + if "ff.2" in orig_key: + orig_key = orig_key.replace("ff.2", "output.dense") + if "ff" in orig_key: + orig_key = orig_key.replace("ff", "output.dense") + if "mlm_class" in orig_key: + orig_key = orig_key.replace("mlm.mlm_class", "cls.predictions.decoder") + if "mlm" in orig_key: + orig_key = orig_key.replace("mlm", "cls.predictions.transform") + if "backbone.backbone.encoders" in orig_key: + orig_key = orig_key.replace("backbone.backbone.encoders", "encoder.layer") + if "cls" not in orig_key: + orig_key = "mra." + orig_key + + return orig_key + + +def convert_checkpoint_helper(max_position_embeddings, orig_state_dict): + for key in orig_state_dict.copy(): + val = orig_state_dict.pop(key) + + if ("pooler" in key) or ("sen_class" in key): + continue + else: + orig_state_dict[rename_key(key)] = val + + orig_state_dict["cls.predictions.bias"] = orig_state_dict["cls.predictions.decoder.bias"] + orig_state_dict["mra.embeddings.position_ids"] = torch.arange(max_position_embeddings).expand((1, -1)) + 2 + + return orig_state_dict + + +def convert_mra_checkpoint(checkpoint_path, mra_config_file, pytorch_dump_path): + orig_state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)["model_state_dict"] + config = MraConfig.from_json_file(mra_config_file) + model = MraForMaskedLM(config) + + new_state_dict = convert_checkpoint_helper(config.max_position_embeddings, orig_state_dict) + + print(model.load_state_dict(new_state_dict)) + model.eval() + model.save_pretrained(pytorch_dump_path) + + print(f"Checkpoint successfully converted. Model saved at {pytorch_dump_path}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + # Required parameters + parser.add_argument( + "--pytorch_model_path", default=None, type=str, required=True, help="Path to Mra pytorch checkpoint." + ) + parser.add_argument( + "--config_file", + default=None, + type=str, + required=True, + help="The json file for Mra model config.", + ) + parser.add_argument( + "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model." + ) + args = parser.parse_args() + convert_mra_checkpoint(args.pytorch_model_path, args.config_file, args.pytorch_dump_path) diff --git a/third_party/transformers/src/transformers/models/mra/modeling_mra.py b/third_party/transformers/src/transformers/models/mra/modeling_mra.py new file mode 100644 index 0000000000000000000000000000000000000000..f36174d09792c45d006b5e10c856aee1d82a64a0 --- /dev/null +++ b/third_party/transformers/src/transformers/models/mra/modeling_mra.py @@ -0,0 +1,1332 @@ +# Copyright 2023 University of Wisconsin-Madison 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 MRA model.""" + +import math + +import torch +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from ... import initialization as init +from ...activations import ACT2FN +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import ( + BaseModelOutputWithCrossAttentions, + MaskedLMOutput, + MultipleChoiceModelOutput, + QuestionAnsweringModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) +from ...modeling_utils import PreTrainedModel +from ...pytorch_utils import apply_chunking_to_forward +from ...utils import ( + auto_docstring, + is_cuda_platform, + is_kernels_available, + is_ninja_available, + is_torch_cuda_available, + logging, +) +from .configuration_mra import MraConfig + + +logger = logging.get_logger(__name__) + +mra_cuda_kernel = None + + +def load_cuda_kernels(): + global mra_cuda_kernel + if not is_kernels_available(): + raise ImportError("kernels is not installed, please install it with `pip install kernels`") + from ...integrations.hub_kernels import get_kernel + + mra_cuda_kernel = get_kernel("kernels-community/mra") + + +def sparse_max(sparse_qk_prod, indices, query_num_block, key_num_block): + """ + Computes maximum values for softmax stability. + """ + if len(sparse_qk_prod.size()) != 4: + raise ValueError("sparse_qk_prod must be a 4-dimensional tensor.") + + if len(indices.size()) != 2: + raise ValueError("indices must be a 2-dimensional tensor.") + + if sparse_qk_prod.size(2) != 32: + raise ValueError("The size of the second dimension of sparse_qk_prod must be 32.") + + if sparse_qk_prod.size(3) != 32: + raise ValueError("The size of the third dimension of sparse_qk_prod must be 32.") + + index_vals = sparse_qk_prod.max(dim=-2).values.transpose(-1, -2) + index_vals = index_vals.contiguous() + + indices = indices.int() + indices = indices.contiguous() + + max_vals, max_vals_scatter = mra_cuda_kernel.index_max(index_vals, indices, query_num_block, key_num_block) + max_vals_scatter = max_vals_scatter.transpose(-1, -2)[:, :, None, :] + + return max_vals, max_vals_scatter + + +def sparse_mask(mask, indices, block_size=32): + """ + Converts attention mask to a sparse mask for high resolution logits. + """ + if len(mask.size()) != 2: + raise ValueError("mask must be a 2-dimensional tensor.") + + if len(indices.size()) != 2: + raise ValueError("indices must be a 2-dimensional tensor.") + + if mask.shape[0] != indices.shape[0]: + raise ValueError("mask and indices must have the same size in the zero-th dimension.") + + batch_size, seq_len = mask.shape + num_block = seq_len // block_size + + batch_idx = torch.arange(indices.size(0), dtype=torch.long, device=indices.device) + mask = mask.reshape(batch_size, num_block, block_size) + mask = mask[batch_idx[:, None], (indices % num_block).long(), :] + + return mask + + +def mm_to_sparse(dense_query, dense_key, indices, block_size=32): + """ + Performs Sampled Dense Matrix Multiplication. + """ + batch_size, query_size, dim = dense_query.size() + _, key_size, dim = dense_key.size() + + if query_size % block_size != 0: + raise ValueError("query_size (size of first dimension of dense_query) must be divisible by block_size.") + + if key_size % block_size != 0: + raise ValueError("key_size (size of first dimension of dense_key) must be divisible by block_size.") + + dense_query = dense_query.reshape(batch_size, query_size // block_size, block_size, dim).transpose(-1, -2) + dense_key = dense_key.reshape(batch_size, key_size // block_size, block_size, dim).transpose(-1, -2) + + if len(dense_query.size()) != 4: + raise ValueError("dense_query must be a 4-dimensional tensor.") + + if len(dense_key.size()) != 4: + raise ValueError("dense_key must be a 4-dimensional tensor.") + + if len(indices.size()) != 2: + raise ValueError("indices must be a 2-dimensional tensor.") + + if dense_query.size(3) != 32: + raise ValueError("The third dimension of dense_query must be 32.") + + if dense_key.size(3) != 32: + raise ValueError("The third dimension of dense_key must be 32.") + + dense_query = dense_query.contiguous() + dense_key = dense_key.contiguous() + + indices = indices.int() + indices = indices.contiguous() + + return mra_cuda_kernel.mm_to_sparse(dense_query, dense_key, indices.int()) + + +def sparse_dense_mm(sparse_query, indices, dense_key, query_num_block, block_size=32): + """ + Performs matrix multiplication of a sparse matrix with a dense matrix. + """ + batch_size, key_size, dim = dense_key.size() + + if key_size % block_size != 0: + raise ValueError("key_size (size of first dimension of dense_key) must be divisible by block_size.") + + if sparse_query.size(2) != block_size: + raise ValueError("The size of the second dimension of sparse_query must be equal to the block_size.") + + if sparse_query.size(3) != block_size: + raise ValueError("The size of the third dimension of sparse_query must be equal to the block_size.") + + dense_key = dense_key.reshape(batch_size, key_size // block_size, block_size, dim).transpose(-1, -2) + + if len(sparse_query.size()) != 4: + raise ValueError("sparse_query must be a 4-dimensional tensor.") + + if len(dense_key.size()) != 4: + raise ValueError("dense_key must be a 4-dimensional tensor.") + + if len(indices.size()) != 2: + raise ValueError("indices must be a 2-dimensional tensor.") + + if dense_key.size(3) != 32: + raise ValueError("The size of the third dimension of dense_key must be 32.") + + sparse_query = sparse_query.contiguous() + + indices = indices.int() + indices = indices.contiguous() + dense_key = dense_key.contiguous() + + dense_qk_prod = mra_cuda_kernel.sparse_dense_mm(sparse_query, indices, dense_key, query_num_block) + dense_qk_prod = dense_qk_prod.transpose(-1, -2).reshape(batch_size, query_num_block * block_size, dim) + return dense_qk_prod + + +def transpose_indices(indices, dim_1_block, dim_2_block): + return ((indices % dim_2_block) * dim_1_block + torch.div(indices, dim_2_block, rounding_mode="floor")).long() + + +class MraSampledDenseMatMul(torch.autograd.Function): + @staticmethod + def forward(ctx, dense_query, dense_key, indices, block_size): + sparse_qk_prod = mm_to_sparse(dense_query, dense_key, indices, block_size) + ctx.save_for_backward(dense_query, dense_key, indices) + ctx.block_size = block_size + return sparse_qk_prod + + @staticmethod + def backward(ctx, grad): + dense_query, dense_key, indices = ctx.saved_tensors + block_size = ctx.block_size + query_num_block = dense_query.size(1) // block_size + key_num_block = dense_key.size(1) // block_size + indices_T = transpose_indices(indices, query_num_block, key_num_block) + grad_key = sparse_dense_mm(grad.transpose(-1, -2), indices_T, dense_query, key_num_block) + grad_query = sparse_dense_mm(grad, indices, dense_key, query_num_block) + return grad_query, grad_key, None, None + + @staticmethod + def operator_call(dense_query, dense_key, indices, block_size=32): + return MraSampledDenseMatMul.apply(dense_query, dense_key, indices, block_size) + + +class MraSparseDenseMatMul(torch.autograd.Function): + @staticmethod + def forward(ctx, sparse_query, indices, dense_key, query_num_block): + sparse_qk_prod = sparse_dense_mm(sparse_query, indices, dense_key, query_num_block) + ctx.save_for_backward(sparse_query, indices, dense_key) + ctx.query_num_block = query_num_block + return sparse_qk_prod + + @staticmethod + def backward(ctx, grad): + sparse_query, indices, dense_key = ctx.saved_tensors + query_num_block = ctx.query_num_block + key_num_block = dense_key.size(1) // sparse_query.size(-1) + indices_T = transpose_indices(indices, query_num_block, key_num_block) + grad_key = sparse_dense_mm(sparse_query.transpose(-1, -2), indices_T, grad, key_num_block) + grad_query = mm_to_sparse(grad, dense_key, indices) + return grad_query, None, grad_key, None + + @staticmethod + def operator_call(sparse_query, indices, dense_key, query_num_block): + return MraSparseDenseMatMul.apply(sparse_query, indices, dense_key, query_num_block) + + +class MraReduceSum: + @staticmethod + def operator_call(sparse_query, indices, query_num_block, key_num_block): + batch_size, num_block, block_size, _ = sparse_query.size() + + if len(sparse_query.size()) != 4: + raise ValueError("sparse_query must be a 4-dimensional tensor.") + + if len(indices.size()) != 2: + raise ValueError("indices must be a 2-dimensional tensor.") + + _, _, block_size, _ = sparse_query.size() + batch_size, num_block = indices.size() + + sparse_query = sparse_query.sum(dim=2).reshape(batch_size * num_block, block_size) + + batch_idx = torch.arange(indices.size(0), dtype=torch.long, device=indices.device) + global_idxes = ( + torch.div(indices, key_num_block, rounding_mode="floor").long() + batch_idx[:, None] * query_num_block + ).reshape(batch_size * num_block) + temp = torch.zeros( + (batch_size * query_num_block, block_size), dtype=sparse_query.dtype, device=sparse_query.device + ) + output = temp.index_add(0, global_idxes, sparse_query).reshape(batch_size, query_num_block, block_size) + + output = output.reshape(batch_size, query_num_block * block_size) + return output + + +def get_low_resolution_logit(query, key, block_size, mask=None, value=None): + """ + Compute low resolution approximation. + """ + batch_size, seq_len, head_dim = query.size() + + num_block_per_row = seq_len // block_size + + value_hat = None + if mask is not None: + token_count = mask.reshape(batch_size, num_block_per_row, block_size).sum(dim=-1) + query_hat = query.reshape(batch_size, num_block_per_row, block_size, head_dim).sum(dim=-2) / ( + token_count[:, :, None] + 1e-6 + ) + key_hat = key.reshape(batch_size, num_block_per_row, block_size, head_dim).sum(dim=-2) / ( + token_count[:, :, None] + 1e-6 + ) + if value is not None: + value_hat = value.reshape(batch_size, num_block_per_row, block_size, head_dim).sum(dim=-2) / ( + token_count[:, :, None] + 1e-6 + ) + else: + token_count = block_size * torch.ones(batch_size, num_block_per_row, dtype=torch.float, device=query.device) + query_hat = query.reshape(batch_size, num_block_per_row, block_size, head_dim).mean(dim=-2) + key_hat = key.reshape(batch_size, num_block_per_row, block_size, head_dim).mean(dim=-2) + if value is not None: + value_hat = value.reshape(batch_size, num_block_per_row, block_size, head_dim).mean(dim=-2) + + low_resolution_logit = torch.matmul(query_hat, key_hat.transpose(-1, -2)) / math.sqrt(head_dim) + + low_resolution_logit_row_max = low_resolution_logit.max(dim=-1, keepdims=True).values + + if mask is not None: + low_resolution_logit = ( + low_resolution_logit - 1e4 * ((token_count[:, None, :] * token_count[:, :, None]) < 0.5).float() + ) + + return low_resolution_logit, token_count, low_resolution_logit_row_max, value_hat + + +def get_block_idxes( + low_resolution_logit, num_blocks, approx_mode, initial_prior_first_n_blocks, initial_prior_diagonal_n_blocks +): + """ + Compute the indices of the subset of components to be used in the approximation. + """ + batch_size, total_blocks_per_row, _ = low_resolution_logit.shape + + if initial_prior_diagonal_n_blocks > 0: + offset = initial_prior_diagonal_n_blocks // 2 + temp_mask = torch.ones(total_blocks_per_row, total_blocks_per_row, device=low_resolution_logit.device) + diagonal_mask = torch.tril(torch.triu(temp_mask, diagonal=-offset), diagonal=offset) + low_resolution_logit = low_resolution_logit + diagonal_mask[None, :, :] * 5e3 + + if initial_prior_first_n_blocks > 0: + low_resolution_logit[:, :initial_prior_first_n_blocks, :] = ( + low_resolution_logit[:, :initial_prior_first_n_blocks, :] + 5e3 + ) + low_resolution_logit[:, :, :initial_prior_first_n_blocks] = ( + low_resolution_logit[:, :, :initial_prior_first_n_blocks] + 5e3 + ) + + top_k_vals = torch.topk( + low_resolution_logit.reshape(batch_size, -1), num_blocks, dim=-1, largest=True, sorted=False + ) + indices = top_k_vals.indices + + if approx_mode == "full": + threshold = top_k_vals.values.min(dim=-1).values + high_resolution_mask = (low_resolution_logit >= threshold[:, None, None]).float() + elif approx_mode == "sparse": + high_resolution_mask = None + else: + raise ValueError(f"{approx_mode} is not a valid approx_model value.") + + return indices, high_resolution_mask + + +def mra2_attention( + query, + key, + value, + mask, + num_blocks, + approx_mode, + block_size=32, + initial_prior_first_n_blocks=0, + initial_prior_diagonal_n_blocks=0, +): + """ + Use Mra to approximate self-attention. + """ + if mra_cuda_kernel is None: + return torch.zeros_like(query).requires_grad_() + + batch_size, num_head, seq_len, head_dim = query.size() + meta_batch = batch_size * num_head + + if seq_len % block_size != 0: + raise ValueError("sequence length must be divisible by the block_size.") + + num_block_per_row = seq_len // block_size + + query = query.reshape(meta_batch, seq_len, head_dim) + key = key.reshape(meta_batch, seq_len, head_dim) + value = value.reshape(meta_batch, seq_len, head_dim) + + if mask is not None: + query = query * mask[:, :, None] + key = key * mask[:, :, None] + value = value * mask[:, :, None] + + if approx_mode == "full": + low_resolution_logit, token_count, low_resolution_logit_row_max, value_hat = get_low_resolution_logit( + query, key, block_size, mask, value + ) + elif approx_mode == "sparse": + with torch.no_grad(): + low_resolution_logit, token_count, low_resolution_logit_row_max, _ = get_low_resolution_logit( + query, key, block_size, mask + ) + else: + raise Exception('approx_mode must be "full" or "sparse"') + + with torch.no_grad(): + low_resolution_logit_normalized = low_resolution_logit - low_resolution_logit_row_max + indices, high_resolution_mask = get_block_idxes( + low_resolution_logit_normalized, + num_blocks, + approx_mode, + initial_prior_first_n_blocks, + initial_prior_diagonal_n_blocks, + ) + + high_resolution_logit = MraSampledDenseMatMul.operator_call( + query, key, indices, block_size=block_size + ) / math.sqrt(head_dim) + max_vals, max_vals_scatter = sparse_max(high_resolution_logit, indices, num_block_per_row, num_block_per_row) + high_resolution_logit = high_resolution_logit - max_vals_scatter + if mask is not None: + high_resolution_logit = high_resolution_logit - 1e4 * (1 - sparse_mask(mask, indices)[:, :, :, None]) + high_resolution_attn = torch.exp(high_resolution_logit) + high_resolution_attn_out = MraSparseDenseMatMul.operator_call( + high_resolution_attn, indices, value, num_block_per_row + ) + high_resolution_normalizer = MraReduceSum.operator_call( + high_resolution_attn, indices, num_block_per_row, num_block_per_row + ) + + if approx_mode == "full": + low_resolution_attn = ( + torch.exp(low_resolution_logit - low_resolution_logit_row_max - 1e4 * high_resolution_mask) + * token_count[:, None, :] + ) + + low_resolution_attn_out = ( + torch.matmul(low_resolution_attn, value_hat)[:, :, None, :] + .repeat(1, 1, block_size, 1) + .reshape(meta_batch, seq_len, head_dim) + ) + low_resolution_normalizer = ( + low_resolution_attn.sum(dim=-1)[:, :, None].repeat(1, 1, block_size).reshape(meta_batch, seq_len) + ) + + log_correction = low_resolution_logit_row_max.repeat(1, 1, block_size).reshape(meta_batch, seq_len) - max_vals + if mask is not None: + log_correction = log_correction * mask + + low_resolution_corr = torch.exp(log_correction * (log_correction <= 0).float()) + low_resolution_attn_out = low_resolution_attn_out * low_resolution_corr[:, :, None] + low_resolution_normalizer = low_resolution_normalizer * low_resolution_corr + + high_resolution_corr = torch.exp(-log_correction * (log_correction > 0).float()) + high_resolution_attn_out = high_resolution_attn_out * high_resolution_corr[:, :, None] + high_resolution_normalizer = high_resolution_normalizer * high_resolution_corr + + context_layer = (high_resolution_attn_out + low_resolution_attn_out) / ( + high_resolution_normalizer[:, :, None] + low_resolution_normalizer[:, :, None] + 1e-6 + ) + + elif approx_mode == "sparse": + context_layer = high_resolution_attn_out / (high_resolution_normalizer[:, :, None] + 1e-6) + else: + raise Exception('config.approx_mode must be "full" or "sparse"') + + if mask is not None: + context_layer = context_layer * mask[:, :, None] + + context_layer = context_layer.reshape(batch_size, num_head, seq_len, head_dim) + + return context_layer + + +class MraEmbeddings(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 + 2, 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)) + 2) + self.register_buffer( + "token_type_ids", + torch.zeros(self.position_ids.size(), dtype=torch.long, device=self.position_ids.device), + 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.dropout(embeddings) + return embeddings + + +class MraSelfAttention(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})" + ) + + kernel_loaded = mra_cuda_kernel is not None + if is_torch_cuda_available() and is_cuda_platform() and is_ninja_available() and not kernel_loaded: + try: + load_cuda_kernels() + except Exception as e: + logger.warning(f"Could not load the custom kernel for multi-scale deformable attention: {e}") + + 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.num_block = (config.max_position_embeddings // 32) * config.block_per_row + self.num_block = min(self.num_block, int((config.max_position_embeddings // 32) ** 2)) + + self.approx_mode = config.approx_mode + self.initial_prior_first_n_blocks = config.initial_prior_first_n_blocks + self.initial_prior_diagonal_n_blocks = config.initial_prior_diagonal_n_blocks + + def forward(self, hidden_states, attention_mask=None): + batch_size, seq_len, _ = 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) + ) + + # revert changes made by get_extended_attention_mask + attention_mask = 1.0 + attention_mask / 10000.0 + attention_mask = ( + attention_mask.squeeze() + .repeat(1, self.num_attention_heads, 1) + .reshape(batch_size * self.num_attention_heads, seq_len) + .int() + ) + + # The CUDA kernels are most efficient with inputs whose size is a multiple of a GPU's warp size (32). Inputs + # smaller than this are padded with zeros. + gpu_warp_size = 32 + + if self.attention_head_size < gpu_warp_size: + pad_size = batch_size, self.num_attention_heads, seq_len, gpu_warp_size - self.attention_head_size + + query_layer = torch.cat([query_layer, torch.zeros(pad_size, device=query_layer.device)], dim=-1) + key_layer = torch.cat([key_layer, torch.zeros(pad_size, device=key_layer.device)], dim=-1) + value_layer = torch.cat([value_layer, torch.zeros(pad_size, device=value_layer.device)], dim=-1) + + context_layer = mra2_attention( + query_layer.float(), + key_layer.float(), + value_layer.float(), + attention_mask.float(), + self.num_block, + approx_mode=self.approx_mode, + initial_prior_first_n_blocks=self.initial_prior_first_n_blocks, + initial_prior_diagonal_n_blocks=self.initial_prior_diagonal_n_blocks, + ) + + if self.attention_head_size < gpu_warp_size: + context_layer = context_layer[:, :, :, : self.attention_head_size] + + context_layer = context_layer.reshape(batch_size, self.num_attention_heads, seq_len, self.attention_head_size) + + 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,) + + return outputs + + +# Copied from transformers.models.bert.modeling_bert.BertSelfOutput +class MraSelfOutput(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 MraAttention(nn.Module): + def __init__(self, config): + super().__init__() + self.self = MraSelfAttention(config) + self.output = MraSelfOutput(config) + + def forward(self, hidden_states, attention_mask=None): + self_outputs = self.self(hidden_states, attention_mask) + 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.bert.modeling_bert.BertIntermediate +class MraIntermediate(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 MraOutput(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 MraLayer(GradientCheckpointingLayer): + def __init__(self, config): + super().__init__() + self.chunk_size_feed_forward = config.chunk_size_feed_forward + self.seq_len_dim = 1 + self.attention = MraAttention(config) + self.add_cross_attention = config.add_cross_attention + self.intermediate = MraIntermediate(config) + self.output = MraOutput(config) + + def forward(self, hidden_states, attention_mask=None): + self_attention_outputs = self.attention(hidden_states, attention_mask) + attention_output = self_attention_outputs[0] + + outputs = self_attention_outputs[1:] # add self attentions if we output attention weights + + layer_output = apply_chunking_to_forward( + self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output + ) + outputs = (layer_output,) + outputs + + return outputs + + 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 MraEncoder(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.layer = nn.ModuleList([MraLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states, + attention_mask=None, + 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, attention_mask) + + 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 BaseModelOutputWithCrossAttentions( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + ) + + +# Copied from transformers.models.bert.modeling_bert.BertPredictionHeadTransform +class MraPredictionHeadTransform(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 + + +# Copied from transformers.models.bert.modeling_bert.BertLMPredictionHead with Bert->Mra +class MraLMPredictionHead(nn.Module): + def __init__(self, config): + super().__init__() + self.transform = MraPredictionHeadTransform(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 + + +# Copied from transformers.models.bert.modeling_bert.BertOnlyMLMHead with Bert->Mra +class MraOnlyMLMHead(nn.Module): + def __init__(self, config): + super().__init__() + self.predictions = MraLMPredictionHead(config) + + def forward(self, sequence_output: torch.Tensor) -> torch.Tensor: + prediction_scores = self.predictions(sequence_output) + return prediction_scores + + +@auto_docstring +# Copied from transformers.models.yoso.modeling_yoso.YosoPreTrainedModel with Yoso->Mra,yoso->mra +class MraPreTrainedModel(PreTrainedModel): + config: MraConfig + base_model_prefix = "mra" + supports_gradient_checkpointing = True + + @torch.no_grad() + def _init_weights(self, module: nn.Module): + """Initialize the weights""" + super()._init_weights(module) + if isinstance(module, MraLMPredictionHead): + init.zeros_(module.bias) + elif isinstance(module, MraEmbeddings): + init.copy_(module.position_ids, torch.arange(module.position_ids.shape[-1]).expand((1, -1)) + 2) + init.zeros_(module.token_type_ids) + + +@auto_docstring +class MraModel(MraPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.config = config + + self.embeddings = MraEmbeddings(config) + self.encoder = MraEncoder(config) + + # 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.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_hidden_states: bool | None = None, + return_dict: bool | None = None, + **kwargs, + ) -> tuple | BaseModelOutputWithCrossAttentions: + 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") + + 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) + + 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) + + # 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, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = encoder_outputs[0] + + if not return_dict: + return (sequence_output,) + encoder_outputs[1:] + + return BaseModelOutputWithCrossAttentions( + last_hidden_state=sequence_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + cross_attentions=encoder_outputs.cross_attentions, + ) + + +@auto_docstring +class MraForMaskedLM(MraPreTrainedModel): + _tied_weights_keys = { + "cls.predictions.decoder.bias": "cls.predictions.bias", + "cls.predictions.decoder.weight": "mra.embeddings.word_embeddings.weight", + } + + def __init__(self, config): + super().__init__(config) + + self.mra = MraModel(config) + self.cls = MraOnlyMLMHead(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, + 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, + 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.mra( + input_ids, + attention_mask=attention_mask, + 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[1:] + 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, + attentions=outputs.attentions, + ) + + +# Copied from transformers.models.yoso.modeling_yoso.YosoClassificationHead with Yoso->Mra +class MraClassificationHead(nn.Module): + """Head for sentence-level classification tasks.""" + + def __init__(self, config): + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.out_proj = nn.Linear(config.hidden_size, config.num_labels) + + self.config = config + + def forward(self, features, **kwargs): + x = features[:, 0, :] # take token (equiv. to [CLS]) + x = self.dropout(x) + x = self.dense(x) + x = ACT2FN[self.config.hidden_act](x) + x = self.dropout(x) + x = self.out_proj(x) + return x + + +@auto_docstring( + custom_intro=""" + MRA 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 MraForSequenceClassification(MraPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + self.mra = MraModel(config) + self.classifier = MraClassificationHead(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, + 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.mra( + input_ids, + attention_mask=attention_mask, + 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.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) + 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 MraForMultipleChoice(MraPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.mra = MraModel(config) + self.pre_classifier = nn.Linear(config.hidden_size, config.hidden_size) + 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, + 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, + 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 + 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.mra( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + hidden_state = outputs[0] # (bs * num_choices, seq_len, dim) + pooled_output = hidden_state[:, 0] # (bs * num_choices, dim) + pooled_output = self.pre_classifier(pooled_output) # (bs * num_choices, dim) + pooled_output = nn.ReLU()(pooled_output) # (bs * num_choices, dim) + 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 MraForTokenClassification(MraPreTrainedModel): + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.mra = MraModel(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, + 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, + 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.mra( + input_ids, + attention_mask=attention_mask, + 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 + 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)) + + 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 MraForQuestionAnswering(MraPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + config.num_labels = 2 + self.num_labels = config.num_labels + + self.mra = MraModel(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, + 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.mra( + input_ids, + attention_mask=attention_mask, + 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) + end_logits = end_logits.squeeze(-1) + + 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[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__ = [ + "MraForMaskedLM", + "MraForMultipleChoice", + "MraForQuestionAnswering", + "MraForSequenceClassification", + "MraForTokenClassification", + "MraLayer", + "MraModel", + "MraPreTrainedModel", +] diff --git a/third_party/transformers/src/transformers/models/nemotron_h/__init__.py b/third_party/transformers/src/transformers/models/nemotron_h/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bf5b9326beeae718ed8a95cb1a3bac078075f94e --- /dev/null +++ b/third_party/transformers/src/transformers/models/nemotron_h/__init__.py @@ -0,0 +1,27 @@ +# 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_nemotron_h import * + from .modeling_nemotron_h 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/nemotron_h/configuration_nemotron_h.py b/third_party/transformers/src/transformers/models/nemotron_h/configuration_nemotron_h.py new file mode 100644 index 0000000000000000000000000000000000000000..4d9361f1e5d2ef3b9a555535b1276567155a330a --- /dev/null +++ b/third_party/transformers/src/transformers/models/nemotron_h/configuration_nemotron_h.py @@ -0,0 +1,274 @@ +# Copyright 2024-2025 NVIDIA Corporation 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. +"""NemotronH 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/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16") +@strict +class NemotronHConfig(PreTrainedConfig): + r""" + layers_block_type (`list`, *optional*): + Explicit list of layer types for each layer. Each element must be one of: "mamba", "attention", or "moe". + The number of layers is determined by the length of this list. + num_logits_to_keep (`int`, *optional*, defaults to 1): + Number of prompt logits to calculate during generation. If `None`, all logits will be calculated. + use_mamba_kernels (`bool`, *optional*, defaults to `True`): + Flag indicating whether or not to use the fast mamba kernels. + ssm_state_size (`int`, *optional*, defaults to 128): + The dimension of the mamba state space latents. + mamba_hidden_act (`str`, *optional*, defaults to `"silu"`): + The non-linear activation function in the Mamba layers. + n_groups (`int`, *optional*, defaults to 8): + Number of groups for the evolution matrices of the Mamba layers. + expand (`int`, *optional*, defaults to 2): + Expanding factor used to determine the intermediate size in the Mamba layers. + use_conv_bias (`bool`, *optional*, defaults to `True`): + Whether or not to use bias in the convolution layer of the Mamba mixer block. + chunk_size (`int`, *optional*, defaults to 128): + Size of the chunks that will comprise the sequence in the Mamba layers. + mamba_ssm_cache_dtype (`str`, *optional*, defaults to `"float32"`): + Data type for Mamba SSM cache states. + moe_shared_expert_intermediate_size (`int`, *optional*, defaults to 7688): + Dimension of the MLP representations in shared experts. + moe_latent_size (`int`, *optional*): + Latent size for MoE expert projections. If `None`, uses `hidden_size`. + moe_shared_expert_overlap (`bool`, *optional*, defaults to `True`): + Whether shared experts overlap with routed experts. + n_group (`int`, *optional*, defaults to 1): + Number of groups for expert routing. + num_nextn_predict_layers (`int`, *optional*, defaults to 0): + Number of additional layers for multi-token prediction. If 0, multi-token prediction is disabled. + mtp_layers_block_type (`list`, *optional*, defaults to `['attention', 'moe']`): + Explicit list of layer types for multi-token prediction layers when `num_nextn_predict_layers` > 0. + use_bias (`bool`, *optional*, defaults to `False`): + Whether to use bias in the model. + residual_in_fp32 (`bool`, *optional*, defaults to `False`): + Whether or not residuals should be in `float32`. + rescale_prenorm_residual (`bool`, *optional*, defaults to `True`): + Whether to rescale the pre-normalization residual connections. + + ```python + >>> from transformers import NemotronHModel, NemotronHConfig + + >>> # Initializing a NemotronH configuration + >>> configuration = NemotronHConfig() + + >>> # Initializing a model (with random weights) from the configuration + >>> model = NemotronHModel(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ``` + """ + + model_type = "nemotron_h" + attribute_map = {"layer_types": "layers_block_type"} + keys_to_ignore_at_inference = ["past_key_values"] + + vocab_size: int = 131072 + hidden_size: int = 4096 + layers_block_type: list[str] | None = None + tie_word_embeddings: bool = False + use_cache: bool = True + num_logits_to_keep: int = 1 + pad_token_id: int | None = 0 + bos_token_id: int | None = 1 + eos_token_id: int | list[int] | None = 2 + num_attention_heads: int = 32 + num_key_value_heads: int = 8 + head_dim: int = 128 + max_position_embeddings: int = 4096 + attention_bias: bool = False + attention_dropout: float | int = 0.0 + sliding_window: int | None = None + intermediate_size: int = 21504 + mlp_hidden_act: str = "relu2" + mlp_bias: bool = False + use_mamba_kernels: bool = True + ssm_state_size: int = 128 + mamba_num_heads: int = 128 + mamba_head_dim: int = 64 + mamba_hidden_act: str = "silu" + n_groups: int = 8 + conv_kernel: int = 4 + expand: int = 2 + time_step_min: float = 0.001 + time_step_max: float = 0.1 + time_step_limit: list[float] | tuple[float, ...] = (0.0, float("inf")) + time_step_floor: float = 1e-4 + use_conv_bias: bool = True + chunk_size: int = 128 + mamba_proj_bias: bool = False + mamba_ssm_cache_dtype: str = "float32" + n_routed_experts: int = 8 + n_shared_experts: int = 1 + moe_intermediate_size: int = 7688 + moe_shared_expert_intermediate_size: int = 7688 + moe_latent_size: int | None = None + moe_shared_expert_overlap: bool = True + num_experts_per_tok: int = 2 + routed_scaling_factor: float | int = 1.0 + n_group: int = 1 + topk_group: int = 1 + norm_topk_prob: bool = True + num_nextn_predict_layers: int = 0 + mtp_layers_block_type: list[str] | None = None + use_bias: bool = False + initializer_range: float = 0.02 + layer_norm_epsilon: float = 1e-5 + residual_in_fp32: bool = False + hidden_dropout: float | int = 0.0 + rescale_prenorm_residual: bool = True + + def __post_init__(self, **kwargs): + # Backward compatibility; configs expect different names for these fields when init + # but they have to be re-names when creating/saving the config. + self.n_groups = kwargs.pop("mamba_n_groups") if "mamba_n_groups" in kwargs else self.n_groups + self.conv_kernel = kwargs.pop("mamba_d_conv") if "mamba_d_conv" in kwargs else self.conv_kernel + self.expand = kwargs.pop("mamba_expand") if "mamba_expand" in kwargs else self.expand + self.time_step_min = kwargs.pop("mamba_dt_min") if "mamba_dt_min" in kwargs else self.time_step_min + self.time_step_max = kwargs.pop("mamba_dt_max") if "mamba_dt_max" in kwargs else self.time_step_max + self.time_step_limit = kwargs.pop("mamba_dt_limit") if "mamba_dt_limit" in kwargs else self.time_step_limit + self.time_step_floor = ( + kwargs.pop("mamba_dt_init_floor") if "mamba_dt_init_floor" in kwargs else self.time_step_floor + ) + self.use_conv_bias = kwargs.pop("mamba_conv_bias") if "mamba_conv_bias" in kwargs else self.use_conv_bias + self.chunk_size = kwargs.pop("mamba_chunk_size") if "mamba_chunk_size" in kwargs else self.chunk_size + + # Backward compatibility: convert hybrid_override_pattern to layers_block_type + # Always pop hybrid_override_pattern from kwargs to prevent it from being set as an attribute + if "hybrid_override_pattern" in kwargs: + pattern = kwargs.pop("hybrid_override_pattern") + if self.layers_block_type is None: + self.layers_block_type = self._pattern_to_list(pattern) + elif self.layers_block_type is None: + # Default layers_block_type if not provided + self.layers_block_type = ["mamba", "moe", "attention", "moe"] + + # Note: num_hidden_layers is deprecated and ignored if layers_block_type is explicitly provided + # It's only kept for backward compatibility when loading old configs + if self.num_hidden_layers is not None: + # Warn if num_hidden_layers is provided but doesn't match layers_block_type + if len(self.layers_block_type) != self.num_hidden_layers: + logger.warning( + f"num_hidden_layers ({self.num_hidden_layers}) is deprecated and doesn't match " + f"layers_block_type length ({len(self.layers_block_type)}). Using layers_block_type length." + ) + + # Backward compatibility: convert mtp_hybrid_override_pattern to mtp_layers_block_type + # Always pop mtp_hybrid_override_pattern from kwargs to prevent it from being set as an attribute + if self.mtp_layers_block_type is None: + self.mtp_layers_block_type = ["attention", "moe"] + + if "mtp_hybrid_override_pattern" in kwargs: + pattern = kwargs.pop("mtp_hybrid_override_pattern") + if self.mtp_layers_block_type == ["attention", "moe"]: + self.mtp_layers_block_type = self._pattern_to_list(pattern) + + # for backward compatibility + if self.num_key_value_heads is None: + self.num_key_value_heads = self.num_attention_heads + + super().__post_init__(**kwargs) + + @staticmethod + def validate_layers_block_type(self): + """ + Validate layers_block_type list. + """ + if not isinstance(self.layers_block_type, list): + raise ValueError( + f"`layers_block_type` must be a list of strings. Got type: {type(self.layers_block_type)}" + ) + + valid_types = {"mamba", "attention", "moe"} + if not all(block_type in valid_types for block_type in self.layers_block_type): + invalid = set(self.layers_block_type) - valid_types + raise ValueError(f"`layers_block_type` contains invalid types: {invalid}. Must be one of: {valid_types}") + + if self.num_nextn_predict_layers > 0: + if self.mtp_layers_block_type is None: + raise ValueError( + "mtp_layers_block_type is required when num_nextn_predict_layers > 0. " + "Please provide an explicit list of layer types for MTP layers. " + "Example: mtp_layers_block_type=['attention', 'moe']" + ) + + if not isinstance(self.mtp_layers_block_type, list): + raise ValueError( + f"`mtp_layers_block_type` must be a list of strings. Got type: {type(self.mtp_layers_block_type)}" + ) + + valid_types = {"mamba", "attention", "moe"} + if not all(block_type in valid_types for block_type in self.mtp_layers_block_type): + invalid = set(self.mtp_layers_block_type) - valid_types + raise ValueError( + f"`mtp_layers_block_type` contains invalid types: {invalid}. Must be one of: {valid_types}" + ) + + @property + def num_hidden_layers(self) -> int: + """ + Number of hidden layers derived from the length of layers_block_type. + This property replaces the deprecated num_hidden_layers parameter. + """ + return len(self.layers_block_type) + + @num_hidden_layers.setter + def num_hidden_layers(self, value): + """ + Setter for backward compatibility when loading configs. + The value is ignored since num_hidden_layers is computed from layers_block_type. + """ + # Ignore the value - num_hidden_layers is always derived from layers_block_type + pass + + @property + def hybrid_override_pattern(self) -> str: + """ + Backward compatibility property. + Returns the pattern string representation of layers_block_type. + """ + return self._list_to_pattern(self.layers_block_type) + + @property + def mtp_hybrid_override_pattern(self) -> str: + """ + Backward compatibility property. + Returns the pattern string representation of mtp_layers_block_type. + """ + return self._list_to_pattern(self.mtp_layers_block_type) + + @staticmethod + def _list_to_pattern(layers_list: list) -> str: + """Convert list of layer types back to pattern string (for backward compatibility).""" + reverse_mapping = {"mamba": "M", "moe": "E", "attention": "*"} + return "".join(reverse_mapping[layer_type] for layer_type in layers_list) + + @staticmethod + def _pattern_to_list(pattern: str) -> list: + """Convert pattern string to list of layer types (for backward compatibility).""" + pattern_mapping = {"M": "mamba", "E": "moe", "*": "attention"} + return [pattern_mapping[char] for char in pattern] + + +__all__ = ["NemotronHConfig"] diff --git a/third_party/transformers/src/transformers/models/nemotron_h/modeling_nemotron_h.py b/third_party/transformers/src/transformers/models/nemotron_h/modeling_nemotron_h.py new file mode 100644 index 0000000000000000000000000000000000000000..9e264e5cfdcccbf0ca61620320fcc48d6e223407 --- /dev/null +++ b/third_party/transformers/src/transformers/models/nemotron_h/modeling_nemotron_h.py @@ -0,0 +1,1212 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/nemotron_h/modular_nemotron_h.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_nemotron_h.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# Copyright (c) 2025, 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. + + +import math +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 ...generation import GenerationMixin +from ...integrations import ( + lazy_load_kernel, + use_experts_implementation, + 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, CausalLMOutputWithPast +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...models.zamba2.modeling_zamba2 import Zamba2RMSNormGated +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 .configuration_nemotron_h import NemotronHConfig + + +logger = logging.get_logger(__name__) + + +# 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 + + +is_fast_path_available = False + + +class NemotronHMamba2Mixer(nn.Module): + """ + Compute ∆, A, B, C, and D the state space parameters and compute the `contextualized_states`. + A, D are input independent (see Mamba paper [1] Section 3.5.2 "Interpretation of A" for why A isn't selective) + ∆, B, C are input-dependent (this is a key difference between Mamba and the linear time invariant S4, + and is why Mamba is called **selective** state spaces) + """ + + def __init__(self, config: NemotronHConfig, layer_idx: int | None = None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.ssm_state_size = config.ssm_state_size + self.conv_kernel_size = config.conv_kernel + self.intermediate_size = config.mamba_num_heads * config.mamba_head_dim + self.layer_idx = layer_idx + self.use_conv_bias = config.use_conv_bias + self.activation = config.mamba_hidden_act + self.act = ACT2FN[config.mamba_hidden_act] + self.use_mem_eff_path = True + + self.n_groups = config.n_groups + self.head_dim = config.mamba_head_dim + self.num_heads = config.mamba_num_heads + self.chunk_size = config.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.use_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=config.use_bias, + ) + # selective projection used to make dt, B and C input dependent + + # 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.norm = Zamba2RMSNormGated( + self.intermediate_size, group_size=self.intermediate_size // self.n_groups, eps=config.layer_norm_epsilon + ) + self.D = nn.Parameter(torch.ones(self.num_heads)) + + self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_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" + ) + + def cuda_kernels_forward( + self, + hidden_states: torch.Tensor, + cache_params: Cache | None = None, + attention_mask: torch.Tensor | None = None, + ): + # set up dimensions for reshapes later + + batch_size, seq_len, _ = hidden_states.shape + groups_time_state_size = self.n_groups * self.ssm_state_size + d_to_remove = 2 * self.intermediate_size + 2 * self.n_groups * self.ssm_state_size + self.num_heads + + # getting projected states from cache if it exists + if cache_params is not None and cache_params.has_previous_state(self.layer_idx): + in_projected_states = self.in_proj(hidden_states.squeeze(1)) # (B 2D) + d_mlp = (in_projected_states.shape[-1] - d_to_remove) // 2 + split_projection_dim = [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads] + _, _, gate, hidden_states_B_C, dt = torch.split(in_projected_states, split_projection_dim, dim=-1) + + 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, + ) + 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=None, + dt_bias=dt_bias, + dt_softplus=True, + ) + hidden_states = hidden_states.view(batch_size, self.num_heads * self.head_dim) + hidden_states = self.norm(hidden_states, gate) + out = self.out_proj(hidden_states)[:, None, ...] + # if no cache is found, calling the kernel + else: + if attention_mask is not None and not torch.all(attention_mask == 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) + # 1. Gated MLP's linear projection + projected_states = self.in_proj(hidden_states) + A = -torch.exp(self.A_log.float()) # (num_heads) or (intermediate_size, state_size) + dt_limit_kwargs = {} if self.time_step_limit is None else {"dt_limit": self.time_step_limit} + if attention_mask is not None: + input_not_masked = torch.all(attention_mask == 1) + else: + input_not_masked = True + + if self.use_mem_eff_path and self.training and cache_params is None and input_not_masked: + out, ssm_state = 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, + activation=self.activation, + rmsnorm_weight=self.norm.weight, + rmsnorm_eps=self.norm.variance_epsilon, + 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=True, + **dt_limit_kwargs, + ) + + else: + gate, hidden_states_B_C, time_step = torch.split( + projected_states, + [self.intermediate_size, self.conv_dim, self.num_heads], + dim=-1, + ) + + # 1D Convolution + if cache_params is not None: + hidden_states_B_C_t = hidden_states_B_C.transpose(1, 2) + conv_state = nn.functional.pad( + hidden_states_B_C_t, (self.conv_kernel_size - hidden_states_B_C_t.shape[-1], 0) + ) + conv_state = cache_params.update_conv_state(conv_state, self.layer_idx) + 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 not torch.all(attention_mask == 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) + 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_bias=self.dt_bias, + dt_softplus=True, + **dt_limit_kwargs, + ) + if ssm_state is not None and cache_params is not None: + 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 + scan_output = self.norm(scan_output, gate) + out = self.out_proj(scan_output) + 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 + # Gated MLP's linear projection + if cache_params is not None and cache_params.has_previous_state(self.layer_idx): + projected_states = self.in_proj(input_states) + else: + if attention_mask is not None: + # tune out hidden states for pad tokens, see https://github.com/state-spaces/mamba/issues/66 + input_states = (input_states * attention_mask[:, :, None]).to(dtype) + projected_states = self.in_proj(input_states) + d_mlp = (projected_states.shape[-1] - 2 * self.intermediate_size - 2 * self.n_groups * self.ssm_state_size- self.num_heads) // 2 + _, _, gate, hidden_states, dt = projected_states.split( + [d_mlp, d_mlp, self.intermediate_size, self.conv_dim, self.num_heads], dim=-1 + ) + hidden_states = hidden_states.transpose(1, 2) + + use_precomputed_state = cache_params is not None and cache_params.has_previous_state(self.layer_idx) + + # Convolution sequence transformation + if use_precomputed_state: + conv_state = cache_params.update_conv_state(hidden_states, self.layer_idx) + hidden_states = torch.sum(conv_state * self.conv1d.weight[:, 0, :], dim=-1) + if self.use_conv_bias: + hidden_states += self.conv1d.bias + hidden_states = self.act(hidden_states).to(dtype)[:, None, ...] # [batch, 1, intermediate_size] : decoding + else: + if cache_params is not None: + conv_state = nn.functional.pad( + hidden_states, + (self.conv_kernel_size - hidden_states.shape[-1], 0) + ) + conv_state = cache_params.update_conv_state(conv_state, self.layer_idx) + + hidden_states = self.act(self.conv1d(hidden_states)[..., :seq_len].transpose(1, 2)) + if attention_mask is not None: + dtype = hidden_states.dtype + # 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) + + hidden_states, B, C = torch.split(hidden_states, [self.intermediate_size, self.n_groups * self.ssm_state_size, self.n_groups * self.ssm_state_size], dim=-1) + A = -torch.exp(self.A_log.float()) # [num_heads] + if use_precomputed_state: + # Note: there is no need to pad parameter matrices here, as there is just one new token + # for batched generation + dt = dt[:, None, ...] if dt.ndim == 2 else 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_min) #, self.time_step_max) + 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) + + # 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] + + # State calculation + ssm_states = cache_params.layers[self.layer_idx].recurrent_states.clone() + ssm_states = ssm_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(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_min) + 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)) + + # First, 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) + + + # Step 2: 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) + + # Step 3: Compute Y_diag (apply to values) + Y_diag = (M[..., None] * hidden_states[:, :, None]).sum(3) + + # (right term of low-rank factorization of off-diagonal blocks; B terms) + + decay_states = torch.exp(A_cumsum[:, :, :, -1:] - A_cumsum) + B_decay_contraction = B * decay_states.permute(0, 2, 3, 1)[..., None] + # permute back B * decay states + states = (B_decay_contraction.permute(0, 1, 3, 2, 4)[..., None] * hidden_states.permute(0, 1, 3, 2, 4)[..., None, :]).sum(dim=3).permute(0, 1, 2, 4, 3) + 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)))) + + states_permuted = states.permute(0, 2, 1, 3, 4) + result = (decay_chunk[..., None, None] * states_permuted[:, :, None, ...]).sum(dim=2) + new_states = result.permute(0, 2, 1, 3, 4) + states, ssm_state = new_states[:, :-1], new_states[:, -1] + + # 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) + # compute Yoff + 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) + if ssm_state is not None and cache_params is not None: + cache_params.update_recurrent_state(ssm_state, self.layer_idx) + + scan_output = self.norm(y, 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(): + # Use cuda stream to avoid NaN when using multiple GPUs, which is caused by multi-GPU synchronization issue. + # Mamba might launch on the default cuda stream that not strictly respect the current Pytorch cuda stream. + # This leads to kernel reading uninitialized memory before the data transfer is complete. + with torch.cuda.stream(torch.cuda.default_stream(hidden_states.device)): + return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask) + + return self.torch_forward(hidden_states, cache_params, attention_mask) + + +@use_kernel_forward_from_hub("RMSNorm") +class NemotronHRMSNorm(nn.Module): + def __init__(self, hidden_size, eps: float = 1e-6) -> None: + """ + NemotronHRMSNorm 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 NemotronHMLP(nn.Module): + def __init__(self, config, intermediate_size=None): + super().__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = intermediate_size or config.intermediate_size + 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.mlp_hidden_act] + + def forward(self, x): + return self.down_proj(self.act_fn(self.up_proj(x))) + + +@use_experts_implementation(has_gate=False) +class NemotronHExperts(nn.Module): + """ + Collection of expert weights stored as 3D tensors. + + **Architecture Note**: Unlike Mixtral or DeepSeek which use gated MLPs, + NemotronH uses a standard MLP architecture with only up_proj and down_proj + """ + + def __init__(self, config): + super().__init__() + self.num_experts = config.n_routed_experts + self.hidden_dim = config.hidden_size + self.intermediate_dim = config.moe_intermediate_size + + # Determine input/output dimension based on whether latent projection is used + input_dim = config.moe_latent_size if config.moe_latent_size is not None else config.hidden_size + + # All expert weights stored as 3D tensors: (num_experts, out_dim, in_dim) + # up_proj: (num_experts, intermediate_dim, input_dim) + self.up_proj = nn.Parameter(torch.empty(self.num_experts, self.intermediate_dim, input_dim)) + # down_proj: (num_experts, input_dim, intermediate_dim) + self.down_proj = nn.Parameter(torch.empty(self.num_experts, input_dim, self.intermediate_dim)) + + self.act_fn = ACT2FN[config.mlp_hidden_act] + + def forward(self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor): + final_hidden_states = torch.zeros_like(hidden_states, dtype=top_k_weights.dtype) + + # Create expert mask to identify which tokens go to which experts + 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) # (num_experts, num_experts_per_tok, num_tokens) + # Only iterate over experts that have at least one token assigned + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero().squeeze(-1) + + for expert_idx in expert_hit: + expert_idx = expert_idx.item() + # Find which tokens are routed to this expert + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + + if token_idx.numel() == 0: + continue + + # Get input for this expert + current_state = hidden_states[token_idx] + + # Expert computation: down_proj(act_fn(up_proj(x))) + # No gating mechanism unlike Mixtral which uses: down_proj(act_fn(gate_proj(x)) * up_proj(x)) + current_hidden_states = torch.nn.functional.linear(current_state, self.up_proj[expert_idx]) + current_hidden_states = self.act_fn(current_hidden_states) + current_hidden_states = torch.nn.functional.linear(current_hidden_states, self.down_proj[expert_idx]) + + # Apply routing weights + current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None] + + # Accumulate into final output + final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype)) + + return final_hidden_states.to(hidden_states.dtype) + + +class NemotronHMoE(nn.Module): + """ + Mixture-of-Experts (MoE) module for NemotronH. + + Unique architectures: + - Uses non-gated MLP experts (NemotronHExperts) instead of gated experts + - Adds optional latent projection for computational efficiency + """ + + def __init__(self, config, layer_idx: int | None = None): + super().__init__() + self.config = config + + # Replace with NemotronH-specific experts (non-gated MLP architecture) + self.experts = NemotronHExperts(config) + self.gate = NemotronHTopkRouter(config) + + # Override shared_experts to use NemotronHMLP with correct intermediate size + self.shared_experts = NemotronHMLP(config=config, intermediate_size=config.moe_shared_expert_intermediate_size) + 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 + + # NemotronH-specific latent projection layers + if config.moe_latent_size is not None: + self.fc1_latent_proj = nn.Linear(config.hidden_size, config.moe_latent_size, bias=config.mlp_bias) + self.fc2_latent_proj = nn.Linear(config.moe_latent_size, config.hidden_size, bias=config.mlp_bias) + else: + self.fc1_latent_proj = nn.Identity() + self.fc2_latent_proj = nn.Identity() + + 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]) + + # NemotronH-specific: latent projection + hidden_states = self.fc1_latent_proj(hidden_states) + hidden_states = self.experts(hidden_states, topk_indices, topk_weights) + hidden_states = self.fc2_latent_proj(hidden_states) + + hidden_states = hidden_states.view(*orig_shape) + hidden_states = hidden_states + self.shared_experts(residuals) + return hidden_states + + +class NemotronHTopkRouter(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.n_routed_experts = config.n_routed_experts + + 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)) + + 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 + + +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 NemotronHAttention(nn.Module): + """Multi-headed attention from 'Attention Is All You Need' paper""" + + def __init__(self, config: NemotronHConfig, 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=False) + self.k_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(config.hidden_size, config.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> 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) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + 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 + + +MIXER_TYPES = { + "mamba": NemotronHMamba2Mixer, + "attention": NemotronHAttention, + "moe": NemotronHMoE, +} + + +class NemotronHBlock(GradientCheckpointingLayer): + """ + A single transformer block in the NemotronH model. + + This block can contain different types of mixers (Mamba, Attention, MLP, or MoE) + depending on the configuration. Each block applies pre-normalization followed by + the mixer, then adds a residual connection. + + Args: + config (`NemotronHConfig`): + Model configuration specifying the block architecture. + layer_idx (`int`): + Index of this block in the model. Used to determine the block type from + `config.layers_block_type[layer_idx]`. + """ + + def __init__(self, config, layer_idx): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.norm = NemotronHRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + + self.block_type = config.layers_block_type[layer_idx] + self.mixer = MIXER_TYPES[self.block_type](config, layer_idx=layer_idx) + + def forward( + self, + hidden_states, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + use_cache: bool | None = False, + **kwargs: Unpack[TransformersKwargs], + ): + residual = hidden_states + hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype)) + + if self.block_type == "mamba": + hidden_states = self.mixer(hidden_states, cache_params=past_key_values, attention_mask=attention_mask) + elif self.block_type == "attention": + hidden_states, _ = self.mixer( + hidden_states=hidden_states, + past_key_values=past_key_values, + attention_mask=attention_mask, + position_ids=position_ids, + user_cache=use_cache, + **kwargs, + ) + else: + hidden_states = self.mixer(hidden_states) + + hidden_states = residual + hidden_states + + return hidden_states + + +class NemotronHPreTrainedModel(PreTrainedModel): + config: NemotronHConfig + base_model_prefix = "model" + _no_split_modules = ["NemotronHBlock"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_flex_attn = True + _is_stateful = True + _can_record_outputs = { + "hidden_states": NemotronHBlock, + "attentions": NemotronHAttention, + } + _keep_in_fp32_modules_strict = [ + "e_score_correction_bias", + ] + _keys_to_ignore_on_load_unexpected = [r"mtp.*"] + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights.""" + super()._init_weights(module) + if isinstance(module, NemotronHMamba2Mixer): + # Initialize A_log and D parameters + A = torch.arange(1, self.config.mamba_num_heads + 1) + init.copy_(module.A_log, torch.log(A)) + init.ones_(module.D) + + dt = torch.exp( + torch.rand(self.config.mamba_num_heads) + * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min)) + + math.log(self.config.time_step_min) + ).clamp(min=self.config.time_step_floor) + + # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + with torch.no_grad(): + init.copy_(module.dt_bias, inv_dt) + module.dt_bias._no_reinit = True + elif isinstance(module, NemotronHTopkRouter): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + init.zeros_(module.e_score_correction_bias) + elif isinstance(module, NemotronHExperts): + # Initialize expert weights + init.normal_(module.up_proj, mean=0.0, std=self.config.initializer_range) + init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) + + if isinstance(module, nn.Linear): + if module.bias is not None: + if not getattr(module.bias, "_no_reinit", False): + init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + init.normal_(module.weight, std=self.config.initializer_range) + + if self.config.rescale_prenorm_residual: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + for name, p in module.named_parameters(): + if name == "out_proj.weight": + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p_new = p / math.sqrt(self.config.num_hidden_layers) + init.copy_(p, p_new) + + +class NemotronHModel(NemotronHPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList([NemotronHBlock(config, layer_idx=idx) for idx in range(config.num_hidden_layers)]) + + self.norm_f = NemotronHRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, new_embeddings): + self.embeddings = new_embeddings + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + input_ids: torch.LongTensor | None = None, + inputs_embeds: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): # ^ is python for xor + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + hidden_states = inputs_embeds + + 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, + input_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) + + # Map block types to their corresponding masks + block_type_to_mask = { + "mamba": mamba_mask, + "attention": causal_mask, + "moe": None, + } + + for layer_idx, mixer_block in enumerate(self.layers): + layer_mask = block_type_to_mask[mixer_block.block_type] + + hidden_states = mixer_block( + hidden_states, + attention_mask=layer_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm_f(hidden_states) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + ) + + 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 + + +# Adapted from transformers.models.jamba.modeling_jamba.JambaForCausalLM with Jamba->NemotronH, JAMBA->NEMOTRON_H +class NemotronHForCausalLM(NemotronHPreTrainedModel, GenerationMixin): + _tied_weights_keys = {} + + def __init__(self, config: NemotronHConfig): + super().__init__(config) + self.model = NemotronHModel(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""" + 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, NemotronHForCausalLM + + >>> model = NemotronHForCausalLM.from_pretrained("Zyphra/NemotronH-7B-v1") + >>> tokenizer = AutoTokenizer.from_pretrained("Zyphra/NemotronH-7B-v1") + + >>> 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, :]).float() + + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.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__ = ["NemotronHPreTrainedModel", "NemotronHModel", "NemotronHForCausalLM"] diff --git a/third_party/transformers/src/transformers/models/nemotron_h/modular_nemotron_h.py b/third_party/transformers/src/transformers/models/nemotron_h/modular_nemotron_h.py new file mode 100644 index 0000000000000000000000000000000000000000..a7433a982f1c982c8d225810fdfe1a06b866c096 --- /dev/null +++ b/third_party/transformers/src/transformers/models/nemotron_h/modular_nemotron_h.py @@ -0,0 +1,521 @@ +# Copyright 2024 The HuggingFace Inc. team. All rights reserved. +# Copyright (c) 2025, 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. + +from __future__ import annotations + +import math + +import torch +from torch import nn + +from ... import initialization as init +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...integrations import use_experts_implementation +from ...masking_utils import create_causal_mask +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from ...modeling_utils import PreTrainedModel +from ...models.deepseek_v3.modeling_deepseek_v3 import DeepseekV3MoE, DeepseekV3TopkRouter +from ...models.jamba.modeling_jamba import JambaAttention +from ...models.llama.modeling_llama import LlamaRMSNorm +from ...models.nemotron.modeling_nemotron import NemotronMLP +from ...models.zamba.modeling_zamba import ZambaForCausalLM +from ...models.zamba2.modeling_zamba2 import Zamba2MambaMixer, Zamba2RMSNormGated +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 capture_outputs +from .configuration_nemotron_h import NemotronHConfig + + +logger = logging.get_logger(__name__) + +is_fast_path_available = False + + +class NemotronHMamba2Mixer(Zamba2MambaMixer): + def __init__(self, config: NemotronHConfig, layer_idx: int | None = None): + super().__init__(config, layer_idx) + self.ssm_state_size = config.ssm_state_size + self.conv_kernel_size = config.conv_kernel + self.intermediate_size = config.mamba_num_heads * config.mamba_head_dim + self.use_conv_bias = config.use_conv_bias + self.activation = config.mamba_hidden_act + self.act = ACT2FN[config.mamba_hidden_act] + self.use_mem_eff_path = True + + self.n_groups = config.n_groups + self.head_dim = config.mamba_head_dim + self.num_heads = config.mamba_num_heads + + self.conv1d = nn.Conv1d( + in_channels=self.conv_dim, + out_channels=self.conv_dim, + bias=config.use_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=config.use_bias, + ) + + self.norm = Zamba2RMSNormGated( + self.intermediate_size, group_size=self.intermediate_size // self.n_groups, eps=config.layer_norm_epsilon + ) + + self.out_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.use_bias) + + 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(): + # Use cuda stream to avoid NaN when using multiple GPUs, which is caused by multi-GPU synchronization issue. + # Mamba might launch on the default cuda stream that not strictly respect the current Pytorch cuda stream. + # This leads to kernel reading uninitialized memory before the data transfer is complete. + with torch.cuda.stream(torch.cuda.default_stream(hidden_states.device)): + return self.cuda_kernels_forward(hidden_states, cache_params, attention_mask) + + return self.torch_forward(hidden_states, cache_params, attention_mask) + + +class NemotronHRMSNorm(LlamaRMSNorm): + pass + + +class NemotronHMLP(NemotronMLP): + def __init__(self, config, intermediate_size=None): + nn.Module.__init__() + self.config = config + self.hidden_size = config.hidden_size + self.intermediate_size = intermediate_size or config.intermediate_size + 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.mlp_hidden_act] + + +@use_experts_implementation(has_gate=False) +class NemotronHExperts(nn.Module): + """ + Collection of expert weights stored as 3D tensors. + + **Architecture Note**: Unlike Mixtral or DeepSeek which use gated MLPs, + NemotronH uses a standard MLP architecture with only up_proj and down_proj + """ + + def __init__(self, config): + super().__init__() + self.num_experts = config.n_routed_experts + self.hidden_dim = config.hidden_size + self.intermediate_dim = config.moe_intermediate_size + + # Determine input/output dimension based on whether latent projection is used + input_dim = config.moe_latent_size if config.moe_latent_size is not None else config.hidden_size + + # All expert weights stored as 3D tensors: (num_experts, out_dim, in_dim) + # up_proj: (num_experts, intermediate_dim, input_dim) + self.up_proj = nn.Parameter(torch.empty(self.num_experts, self.intermediate_dim, input_dim)) + # down_proj: (num_experts, input_dim, intermediate_dim) + self.down_proj = nn.Parameter(torch.empty(self.num_experts, input_dim, self.intermediate_dim)) + + self.act_fn = ACT2FN[config.mlp_hidden_act] + + def forward(self, hidden_states: torch.Tensor, top_k_index: torch.Tensor, top_k_weights: torch.Tensor): + final_hidden_states = torch.zeros_like(hidden_states, dtype=top_k_weights.dtype) + + # Create expert mask to identify which tokens go to which experts + 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) # (num_experts, num_experts_per_tok, num_tokens) + # Only iterate over experts that have at least one token assigned + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero().squeeze(-1) + + for expert_idx in expert_hit: + expert_idx = expert_idx.item() + # Find which tokens are routed to this expert + top_k_pos, token_idx = torch.where(expert_mask[expert_idx]) + + if token_idx.numel() == 0: + continue + + # Get input for this expert + current_state = hidden_states[token_idx] + + # Expert computation: down_proj(act_fn(up_proj(x))) + # No gating mechanism unlike Mixtral which uses: down_proj(act_fn(gate_proj(x)) * up_proj(x)) + current_hidden_states = torch.nn.functional.linear(current_state, self.up_proj[expert_idx]) + current_hidden_states = self.act_fn(current_hidden_states) + current_hidden_states = torch.nn.functional.linear(current_hidden_states, self.down_proj[expert_idx]) + + # Apply routing weights + current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None] + + # Accumulate into final output + final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype)) + + return final_hidden_states.to(hidden_states.dtype) + + +class NemotronHMoE(DeepseekV3MoE): + """ + Mixture-of-Experts (MoE) module for NemotronH. + + Unique architectures: + - Uses non-gated MLP experts (NemotronHExperts) instead of gated experts + - Adds optional latent projection for computational efficiency + """ + + def __init__(self, config, layer_idx: int | None = None): + super().__init__(config) + + # Replace with NemotronH-specific experts (non-gated MLP architecture) + self.experts = NemotronHExperts(config) + self.gate = NemotronHTopkRouter(config) + + # Override shared_experts to use NemotronHMLP with correct intermediate size + self.shared_experts = NemotronHMLP(config=config, intermediate_size=config.moe_shared_expert_intermediate_size) + + # NemotronH-specific latent projection layers + if config.moe_latent_size is not None: + self.fc1_latent_proj = nn.Linear(config.hidden_size, config.moe_latent_size, bias=config.mlp_bias) + self.fc2_latent_proj = nn.Linear(config.moe_latent_size, config.hidden_size, bias=config.mlp_bias) + else: + self.fc1_latent_proj = nn.Identity() + self.fc2_latent_proj = nn.Identity() + + 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]) + + # NemotronH-specific: latent projection + hidden_states = self.fc1_latent_proj(hidden_states) + hidden_states = self.experts(hidden_states, topk_indices, topk_weights) + hidden_states = self.fc2_latent_proj(hidden_states) + + hidden_states = hidden_states.view(*orig_shape) + hidden_states = hidden_states + self.shared_experts(residuals) + return hidden_states + + +class NemotronHTopkRouter(DeepseekV3TopkRouter): + pass + + +class NemotronHAttention(JambaAttention): + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Cache | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + return super().forward(hidden_states, attention_mask, past_key_values, **kwargs) + + +MIXER_TYPES = { + "mamba": NemotronHMamba2Mixer, + "attention": NemotronHAttention, + "moe": NemotronHMoE, +} + + +class NemotronHBlock(GradientCheckpointingLayer): + """ + A single transformer block in the NemotronH model. + + This block can contain different types of mixers (Mamba, Attention, MLP, or MoE) + depending on the configuration. Each block applies pre-normalization followed by + the mixer, then adds a residual connection. + + Args: + config (`NemotronHConfig`): + Model configuration specifying the block architecture. + layer_idx (`int`): + Index of this block in the model. Used to determine the block type from + `config.layers_block_type[layer_idx]`. + """ + + def __init__(self, config, layer_idx): + super().__init__() + self.config = config + self.layer_idx = layer_idx + self.norm = NemotronHRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + + self.block_type = config.layers_block_type[layer_idx] + self.mixer = MIXER_TYPES[self.block_type](config, layer_idx=layer_idx) + + def forward( + self, + hidden_states, + past_key_values: Cache | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + use_cache: bool | None = False, + **kwargs: Unpack[TransformersKwargs], + ): + residual = hidden_states + hidden_states = self.norm(hidden_states.to(dtype=self.norm.weight.dtype)) + + if self.block_type == "mamba": + hidden_states = self.mixer(hidden_states, cache_params=past_key_values, attention_mask=attention_mask) + elif self.block_type == "attention": + hidden_states, _ = self.mixer( + hidden_states=hidden_states, + past_key_values=past_key_values, + attention_mask=attention_mask, + position_ids=position_ids, + user_cache=use_cache, + **kwargs, + ) + else: + hidden_states = self.mixer(hidden_states) + + hidden_states = residual + hidden_states + + return hidden_states + + +class NemotronHPreTrainedModel(PreTrainedModel): + config: NemotronHConfig + base_model_prefix = "model" + _no_split_modules = ["NemotronHBlock"] + _skip_keys_device_placement = ["past_key_values"] + _supports_flash_attn = True + _supports_flash_attn_2 = True + _supports_sdpa = True + _supports_flex_attn = True + _is_stateful = True + _can_record_outputs = { + "hidden_states": NemotronHBlock, + "attentions": NemotronHAttention, + } + _keep_in_fp32_modules_strict = [ + "e_score_correction_bias", + ] + _keys_to_ignore_on_load_unexpected = [r"mtp.*"] + + @torch.no_grad() + def _init_weights(self, module): + """Initialize the weights.""" + super()._init_weights(module) + if isinstance(module, NemotronHMamba2Mixer): + # Initialize A_log and D parameters + A = torch.arange(1, self.config.mamba_num_heads + 1) + init.copy_(module.A_log, torch.log(A)) + init.ones_(module.D) + + dt = torch.exp( + torch.rand(self.config.mamba_num_heads) + * (math.log(self.config.time_step_max) - math.log(self.config.time_step_min)) + + math.log(self.config.time_step_min) + ).clamp(min=self.config.time_step_floor) + + # # Inverse of softplus: https://github.com/pytorch/pytorch/issues/72759 + inv_dt = dt + torch.log(-torch.expm1(-dt)) + with torch.no_grad(): + init.copy_(module.dt_bias, inv_dt) + module.dt_bias._no_reinit = True + elif isinstance(module, NemotronHTopkRouter): + init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + init.zeros_(module.e_score_correction_bias) + elif isinstance(module, NemotronHExperts): + # Initialize expert weights + init.normal_(module.up_proj, mean=0.0, std=self.config.initializer_range) + init.normal_(module.down_proj, mean=0.0, std=self.config.initializer_range) + + if isinstance(module, nn.Linear): + if module.bias is not None: + if not getattr(module.bias, "_no_reinit", False): + init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + init.normal_(module.weight, std=self.config.initializer_range) + + if self.config.rescale_prenorm_residual: + # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme: + # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale + # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers. + # > -- GPT-2 :: https://openai.com/blog/better-language-models/ + # + # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py + for name, p in module.named_parameters(): + if name == "out_proj.weight": + # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block + # Following Pytorch init, except scale by 1/sqrt(2 * n_layer) + # We need to reinit p since this code could be called multiple times + # Having just p *= scale would repeatedly scale it down + init.kaiming_uniform_(p, a=math.sqrt(5)) + with torch.no_grad(): + p_new = p / math.sqrt(self.config.num_hidden_layers) + init.copy_(p, p_new) + + +class NemotronHModel(NemotronHPreTrainedModel): + def __init__(self, config): + super().__init__(config) + + self.embeddings = nn.Embedding(config.vocab_size, config.hidden_size) + self.layers = nn.ModuleList([NemotronHBlock(config, layer_idx=idx) for idx in range(config.num_hidden_layers)]) + + self.norm_f = NemotronHRMSNorm(config.hidden_size, eps=config.layer_norm_epsilon) + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, new_embeddings): + self.embeddings = new_embeddings + + @merge_with_config_defaults + @capture_outputs + def forward( + self, + input_ids: torch.LongTensor | None = None, + inputs_embeds: torch.LongTensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + use_cache: bool | None = None, + attention_mask: torch.Tensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): # ^ is python for xor + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embeddings(input_ids) + + if use_cache and past_key_values is None: + past_key_values = DynamicCache(config=self.config) + + hidden_states = inputs_embeds + + 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, + input_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) + + # Map block types to their corresponding masks + block_type_to_mask = { + "mamba": mamba_mask, + "attention": causal_mask, + "moe": None, + } + + for layer_idx, mixer_block in enumerate(self.layers): + layer_mask = block_type_to_mask[mixer_block.block_type] + + hidden_states = mixer_block( + hidden_states, + attention_mask=layer_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + hidden_states = self.norm_f(hidden_states) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values if use_cache else None, + ) + + 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 NemotronHForCausalLM(ZambaForCausalLM): + _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, + 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: + 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, :]).float() + + loss = None + if labels is not None: + loss = self.loss_function(logits, labels, self.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__ = [ + "NemotronHPreTrainedModel", + "NemotronHModel", + "NemotronHForCausalLM", +] diff --git a/third_party/transformers/src/transformers/models/phobert/__init__.py b/third_party/transformers/src/transformers/models/phobert/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5c14d34aff61362159e0b45e5895ceb08850faf5 --- /dev/null +++ b/third_party/transformers/src/transformers/models/phobert/__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_phobert 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/phobert/tokenization_phobert.py b/third_party/transformers/src/transformers/models/phobert/tokenization_phobert.py new file mode 100644 index 0000000000000000000000000000000000000000..6dd66b7da72d507405e49d1ec1534e160b6215a7 --- /dev/null +++ b/third_party/transformers/src/transformers/models/phobert/tokenization_phobert.py @@ -0,0 +1,349 @@ +# 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 PhoBERT""" + +import os +import re +from shutil import copyfile + +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 PhobertTokenizer(PreTrainedTokenizer): + """ + Construct a PhoBERT 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. + bos_token (`st`, *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, + bos_token="", + eos_token="", + sep_token="", + cls_token="", + unk_token="", + pad_token="", + mask_token="", + **kwargs, + ): + 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 = {} + + super().__init__( + bos_token=bos_token, + eos_token=eos_token, + unk_token=unk_token, + sep_token=sep_token, + cls_token=cls_token, + pad_token=pad_token, + mask_token=mask_token, + **kwargs, + ) + + 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 PhoBERT 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.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 + 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 None: + return [1] + ([0] * len(token_ids_0)) + [1] + return [1] + ([0] * len(token_ids_0)) + [1, 1] + ([0] * len(token_ids_1)) + [1] + + def create_token_type_ids_from_sequences( + self, token_ids_0: list[int], token_ids_1: list[int] | None = None + ) -> list[int]: + """ + Create a mask from the two sequences passed to be used in a sequence-pair classification task. PhoBERT does not + make use of token type ids, therefore a list of zeros is returned. + + Args: + token_ids_0 (`list[int]`): + List of IDs. + token_ids_1 (`list[int]`, *optional*): + Optional second list of IDs for sequence pairs. + + Returns: + `list[int]`: List of zeros. + """ + + sep = [self.sep_token_id] + cls = [self.cls_token_id] + + if token_ids_1 is None: + return len(cls + token_ids_0 + sep) * [0] + return len(cls + token_ids_0 + sep + sep + token_ids_1 + sep) * [0] + + @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 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 + out_vocab_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"] + ) + out_merge_file = os.path.join( + save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["merges_file"] + ) + + if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file): + copyfile(self.vocab_file, out_vocab_file) + elif not os.path.isfile(self.vocab_file): + with open(out_vocab_file, "wb") as fi: + content_spiece_model = self.sp_model.serialized_model_proto() + fi.write(content_spiece_model) + + if os.path.abspath(self.merges_file) != os.path.abspath(out_merge_file): + copyfile(self.merges_file, out_merge_file) + + return out_vocab_file, out_merge_file + + # 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 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) + + +__all__ = ["PhobertTokenizer"] diff --git a/third_party/transformers/src/transformers/models/pp_chart2table/__init__.py b/third_party/transformers/src/transformers/models/pp_chart2table/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..961039282748414733d7d5a107b19ad1735ed7ed --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_chart2table/__init__.py @@ -0,0 +1,30 @@ +# 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_pp_chart2table import * + from .image_processing_pil_pp_chart2table import * + from .image_processing_pp_chart2table import * + from .processing_pp_chart2table 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/pp_chart2table/configuration_pp_chart2table.py b/third_party/transformers/src/transformers/models/pp_chart2table/configuration_pp_chart2table.py new file mode 100644 index 0000000000000000000000000000000000000000..61794f02b5a302f09ac18483f6c49f79a3df9a6e --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_chart2table/configuration_pp_chart2table.py @@ -0,0 +1,131 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/pp_chart2table/modular_pp_chart2table.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_pp_chart2table.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The PaddlePaddle 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. + +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="facebook/sam-vit-huge") +@strict +class PPChart2TableVisionConfig(PreTrainedConfig): + r""" + output_channels (`int`, *optional*, defaults to 256): + Dimensionality of the output channels in the Patch Encoder. + use_abs_pos (`bool`, *optional*, defaults to `True`): + Whether to use absolute position embedding. + use_rel_pos (`bool`, *optional*, defaults to `True`): + Whether to use relative position embedding. + window_size (`int`, *optional*, defaults to 14): + Window size for relative position. + global_attn_indexes (`list[int]`, *optional*, defaults to `[2, 5, 8, 11]`): + The indexes of the global attention layers. + mlp_dim (`int`, *optional*, defaults to 3072): + The dimensionality of the MLP layer in the Transformer encoder. + """ + + base_config_key = "vision_config" + hidden_size: int = 768 + output_channels: int = 256 + num_hidden_layers: int = 12 + num_attention_heads: int = 12 + num_channels: int = 3 + image_size: int | list[int] | tuple[int, int] = 1024 + patch_size: int | list[int] | tuple[int, int] = 16 + hidden_act: str = "gelu" + layer_norm_eps: float = 1e-06 + attention_dropout: float | int = 0.0 + initializer_range: float = 1e-10 + qkv_bias: bool = True + use_abs_pos: bool = True + use_rel_pos: bool = True + window_size: int = 14 + global_attn_indexes: list[int] | tuple[int, ...] = (2, 5, 8, 11) + mlp_dim: int = 3072 + + +@auto_docstring(checkpoint="PaddlePaddle/PP-Chart2Table_safetensors") +@strict +class PPChart2TableConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers import GotOcr2ForConditionalGeneration, PPChart2TableConfig + + >>> # Initializing a PPChart2Table style configuration + >>> configuration = PPChart2TableConfig() + + >>> # Initializing a model from the PaddlePaddle/PP-Chart2Table_safetensors style configuration + >>> model = GotOcr2ForConditionalGeneration(configuration) # underlying architecture is Got Ocr 2 + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "pp_chart2table" + attribute_map = { + "image_token_id": "image_token_index", + } + sub_configs = {"text_config": AutoConfig, "vision_config": PPChart2TableVisionConfig} + + vision_config: dict | PreTrainedConfig | None = None + text_config: dict | PreTrainedConfig | None = None + image_token_index: int = 151859 + image_seq_length: int = 576 + tie_word_embeddings: bool = True + + def __post_init__(self, **kwargs): + if self.vision_config is None: + self.vision_config = PPChart2TableVisionConfig() + elif isinstance(self.vision_config, dict): + self.vision_config = PPChart2TableVisionConfig(**self.vision_config) + + 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"]( + vocab_size=151860, + hidden_size=1024, + intermediate_size=2816, + num_hidden_layers=24, + num_attention_heads=16, + num_key_value_heads=16, + hidden_act="silu", + max_position_embeddings=32768, + initializer_range=0.02, + rms_norm_eps=1e-6, + use_cache=True, + tie_word_embeddings=self.tie_word_embeddings, + rope_theta=1000000.0, + rope_parameters=None, + use_sliding_window=False, + sliding_window=4096, + max_window_layers=21, + attention_dropout=0.0, + ) + + super().__post_init__(**kwargs) + + +__all__ = ["PPChart2TableConfig"] diff --git a/third_party/transformers/src/transformers/models/pp_chart2table/image_processing_pil_pp_chart2table.py b/third_party/transformers/src/transformers/models/pp_chart2table/image_processing_pil_pp_chart2table.py new file mode 100644 index 0000000000000000000000000000000000000000..40cce468b5dcf07b662c45e802edd4fd125f3bcf --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_chart2table/image_processing_pil_pp_chart2table.py @@ -0,0 +1,36 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/pp_chart2table/modular_pp_chart2table.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_pp_chart2table.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The PaddlePaddle 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. + +from ...image_processing_backends import PilBackend +from ...utils import auto_docstring + + +@auto_docstring +class PPChart2TableImageProcessorPil(PilBackend): + resample = 3 + image_mean = [0.48145466, 0.4578275, 0.40821073] + image_std = [0.26862954, 0.26130258, 0.27577711] + size = {"height": 1024, "width": 1024} + do_resize = True + do_rescale = True + do_normalize = True + + +__all__ = ["PPChart2TableImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/pp_chart2table/image_processing_pp_chart2table.py b/third_party/transformers/src/transformers/models/pp_chart2table/image_processing_pp_chart2table.py new file mode 100644 index 0000000000000000000000000000000000000000..b38027aecef96f5e0562f083bf878239f1555796 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_chart2table/image_processing_pp_chart2table.py @@ -0,0 +1,36 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/pp_chart2table/modular_pp_chart2table.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_pp_chart2table.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The PaddlePaddle 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. + +from ...image_processing_backends import TorchvisionBackend +from ...utils import auto_docstring + + +@auto_docstring +class PPChart2TableImageProcessor(TorchvisionBackend): + resample = 3 + image_mean = [0.48145466, 0.4578275, 0.40821073] + image_std = [0.26862954, 0.26130258, 0.27577711] + size = {"height": 1024, "width": 1024} + do_resize = True + do_rescale = True + do_normalize = True + + +__all__ = ["PPChart2TableImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/pp_chart2table/modular_pp_chart2table.py b/third_party/transformers/src/transformers/models/pp_chart2table/modular_pp_chart2table.py new file mode 100644 index 0000000000000000000000000000000000000000..d4953b067fe751102e42412723f8ea2842b72314 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_chart2table/modular_pp_chart2table.py @@ -0,0 +1,94 @@ +# Copyright 2026 The PaddlePaddle 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. + +from huggingface_hub.dataclasses import strict + +from ...feature_extraction_utils import BatchFeature +from ...image_processing_backends import PilBackend, TorchvisionBackend +from ...image_utils import ImageInput +from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import auto_docstring, logging +from ..got_ocr2.configuration_got_ocr2 import GotOcr2Config + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="PaddlePaddle/PP-Chart2Table_safetensors") +@strict +class PPChart2TableConfig(GotOcr2Config): + model_type = "pp_chart2table" + + r""" + Example: + + ```python + >>> from transformers import GotOcr2ForConditionalGeneration, PPChart2TableConfig + + >>> # Initializing a PPChart2Table style configuration + >>> configuration = PPChart2TableConfig() + + >>> # Initializing a model from the PaddlePaddle/PP-Chart2Table_safetensors style configuration + >>> model = GotOcr2ForConditionalGeneration(configuration) # underlying architecture is Got Ocr 2 + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + +@auto_docstring +class PPChart2TableImageProcessor(TorchvisionBackend): + resample = 3 + image_mean = [0.48145466, 0.4578275, 0.40821073] + image_std = [0.26862954, 0.26130258, 0.27577711] + size = {"height": 1024, "width": 1024} + do_resize = True + do_rescale = True + do_normalize = True + + +@auto_docstring +class PPChart2TableImageProcessorPil(PilBackend): + resample = 3 + image_mean = [0.48145466, 0.4578275, 0.40821073] + image_std = [0.26862954, 0.26130258, 0.27577711] + size = {"height": 1024, "width": 1024} + do_resize = True + do_rescale = True + do_normalize = True + + +@auto_docstring +class PPChart2TableProcessor(ProcessorMixin): + def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs): + super().__init__(image_processor, tokenizer, chat_template=chat_template) + + def __call__( + self, + images: ImageInput = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + **kwargs: Unpack[ProcessingKwargs], + ) -> BatchFeature: + if text is None or images is None: + raise ValueError("Both `images` and `text` must be provided") + return super().__call__(images=images, text=text, **kwargs) + + +__all__ = [ + "PPChart2TableConfig", + "PPChart2TableImageProcessor", + "PPChart2TableImageProcessorPil", + "PPChart2TableProcessor", +] diff --git a/third_party/transformers/src/transformers/models/pp_chart2table/processing_pp_chart2table.py b/third_party/transformers/src/transformers/models/pp_chart2table/processing_pp_chart2table.py new file mode 100644 index 0000000000000000000000000000000000000000..6f5e4554b731d7564eaa07751a440d02e8529b44 --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_chart2table/processing_pp_chart2table.py @@ -0,0 +1,45 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/pp_chart2table/modular_pp_chart2table.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_pp_chart2table.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The PaddlePaddle 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. + + +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput +from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack +from ...tokenization_utils_base import PreTokenizedInput, TextInput +from ...utils import auto_docstring + + +@auto_docstring +class PPChart2TableProcessor(ProcessorMixin): + def __init__(self, image_processor=None, tokenizer=None, chat_template=None, **kwargs): + super().__init__(image_processor, tokenizer, chat_template=chat_template) + + def __call__( + self, + images: ImageInput = None, + text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput] = None, + **kwargs: Unpack[ProcessingKwargs], + ) -> BatchFeature: + if text is None or images is None: + raise ValueError("Both `images` and `text` must be provided") + return super().__call__(images=images, text=text, **kwargs) + + +__all__ = ["PPChart2TableProcessor"] diff --git a/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/configuration_pp_ocrv5_server_rec.py b/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/configuration_pp_ocrv5_server_rec.py new file mode 100644 index 0000000000000000000000000000000000000000..b5bffda9c5d75a0f5d2ac8628c5608b9abcd361f --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/configuration_pp_ocrv5_server_rec.py @@ -0,0 +1,75 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/pp_ocrv5_server_rec/modular_pp_ocrv5_server_rec.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_pp_ocrv5_server_rec.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The PaddlePaddle 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. + +from huggingface_hub.dataclasses import strict + +from ...backbone_utils import consolidate_backbone_kwargs_to_config +from ...configuration_utils import PreTrainedConfig +from ...utils import auto_docstring +from ..auto import AutoConfig + + +@auto_docstring(checkpoint="PaddlePaddle/PP-OCRv5_server_rec_safetensors") +@strict +class PPOCRV5ServerRecConfig(PreTrainedConfig): + r""" + head_out_channels (`int`, *optional*, defaults to 18385): + The number of output channels from the PPOCRV5ServerRecHead, responsible for final classification. + """ + + model_type = "pp_ocrv5_server_rec" + sub_configs = {"backbone_config": AutoConfig} + + hidden_act: str = "silu" + backbone_config: dict | PreTrainedConfig | None = None + hidden_size: int = 120 + mlp_ratio: float = 2.0 + depth: int = 2 + head_out_channels: int = 18385 + conv_kernel_size: list | None = None + qkv_bias: bool = True + num_attention_heads: int = 8 + attention_dropout: float | int = 0.0 + layer_norm_eps: float = 1e-6 + + def __post_init__(self, **kwargs): + if self.conv_kernel_size is None: + self.conv_kernel_size = [1, 3] + self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config( + backbone_config=self.backbone_config, + default_config_type="hgnet_v2", + default_config_kwargs={ + "arch": "L", + "return_idx": [0, 1, 2, 3], + "freeze_stem_only": True, + "freeze_at": 0, + "freeze_norm": True, + "lr_mult_list": [1.0, 1.0, 1.0, 1.0, 1.0], + "out_features": ["stage1", "stage2", "stage3", "stage4"], + "stage_downsample": [True, True, True, True], + "stem_strides": [2, 1, 1, 1, 1], + "stage_downsample_strides": [[2, 1], [1, 2], [2, 1], [2, 1]], + }, + **kwargs, + ) + super().__post_init__(**kwargs) + + +__all__ = ["PPOCRV5ServerRecConfig"] diff --git a/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/image_processing_pp_ocrv5_server_rec.py b/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/image_processing_pp_ocrv5_server_rec.py new file mode 100644 index 0000000000000000000000000000000000000000..663788f84fa45fda4df943b7541b1934bbef570f --- /dev/null +++ b/third_party/transformers/src/transformers/models/pp_ocrv5_server_rec/image_processing_pp_ocrv5_server_rec.py @@ -0,0 +1,188 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/pp_ocrv5_server_rec/modular_pp_ocrv5_server_rec.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_pp_ocrv5_server_rec.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2026 The PaddlePaddle 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. + +import math + +import torch +import torchvision.transforms.v2.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 PILImageResampling, SizeDict +from ...processing_utils import ImagesKwargs +from ...utils import auto_docstring, requires_backends +from ...utils.constants import IMAGENET_STANDARD_MEAN, IMAGENET_STANDARD_STD +from ...utils.generic import TensorType + + +class PPOCRV5ServerRecImageProcessorKwargs(ImagesKwargs, total=False): + r""" + max_image_width (`int`, *optional*, defaults to `3200`): + Maximum image width used during resizing. + character_list (`list`, *optional*, defaults to `[]`): + Vocabulary list used for text recognition. + """ + + max_image_width: int + character_list: str + + +@auto_docstring +class PPOCRV5ServerRecImageProcessor(TorchvisionBackend): + resample = PILImageResampling.BILINEAR + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"height": 48, "width": 320} + pad_size = {"height": 48, "width": 320} + do_resize = True + do_rescale = True + do_convert_rgb = True + do_normalize = True + do_pad = True + max_image_width = 3200 + character_list = [] + valid_kwargs = PPOCRV5ServerRecImageProcessorKwargs + + 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, + **kwargs, + ) -> BatchFeature: + # Group images by size for batched resizing + grouped_images, grouped_images_index = group_images_by_shape(images, disable_grouping=disable_grouping) + resized_images_grouped = {} + + # [Key Change] Use get_target_size to calculate target_size for resizing. + shape_list = list(grouped_images.keys()) + target_size = self.get_target_size(shape_list) + + for shape, stacked_images in grouped_images.items(): + if do_resize: + stacked_images = self.resize(image=stacked_images.float(), size=target_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 + ) + processed_images_grouped[shape] = stacked_images + processed_images = reorder_images(processed_images_grouped, grouped_images_index) + + if do_pad and target_size.width < pad_size.width: + processed_images = self.pad(processed_images, pad_size=pad_size, disable_grouping=disable_grouping) + + return BatchFeature(data={"pixel_values": processed_images}, tensor_type=return_tensors) + + def get_target_size(self, shape_list: list[torch.Size]): + """ + Calculate the width and height from the widest image in the batch. + """ + max_width = -1 + max_height = -1 + for height, width in shape_list: + # We need the width and height of the widest image in the batch + if width > max_width: + max_width = width + max_height = height + + default_height, default_width = self.size["height"], self.size["width"] + ratio = max(max_width / max_height, default_width / default_height) + + target_width = int(default_height * ratio) + target_height = default_height + + if target_width > self.max_image_width: + target_width = self.max_image_width + else: + ratio = max_width / float(max_height) + if target_width >= math.ceil(default_height * ratio): + target_width = int(math.ceil(default_height * ratio)) + + return SizeDict(height=target_height, width=target_width) + + def post_process_text_recognition( + self, + predictions, + ) -> list[dict[str, str | float]]: + """ + Post-processes raw model logits to decode the recognized text and its confidence score. + + Args: + predictions: Model outputs with `logits` attribute (probability maps of shape `(batch_size, height, vocab_size)`). + + Returns: + A list of dictionaries, where each dictionary corresponds to an image in the batch. + Each dictionary contains: + - "text" (str): The decoded text string. + - "score" (float): The average confidence score of the characters in the decoded text. + """ + requires_backends(self, ["torch"]) + + logits = predictions.last_hidden_state + batch_size = logits.shape[0] + + preds_prob, preds_idx = logits.max(dim=-1) + results = [] + for idx in range(batch_size): + selection = torch.ones(len(preds_idx[idx]), dtype=torch.bool, device=preds_idx.device) + + # remove_duplicate + selection[1:] = preds_idx[idx][1:] != preds_idx[idx][:-1] + # ignore blank token + selection &= preds_idx[idx] != 0 + + character_list = [] + for text_id in preds_idx[idx][selection]: + character_list.append(self.character_list[text_id]) + + results.append( + { + "text": "".join(character_list), + "score": preds_prob[idx][selection].mean().item(), + } + ) + + return results + + +__all__ = ["PPOCRV5ServerRecImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/smolvlm/__init__.py b/third_party/transformers/src/transformers/models/smolvlm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0ddab5b2b0dbd89521fcadfcd299bd835b4deb1d --- /dev/null +++ b/third_party/transformers/src/transformers/models/smolvlm/__init__.py @@ -0,0 +1,30 @@ +# 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_smolvlm import * + from .image_processing_pil_smolvlm import * + from .image_processing_smolvlm import * + from .modeling_smolvlm import * + from .processing_smolvlm 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/smolvlm/configuration_smolvlm.py b/third_party/transformers/src/transformers/models/smolvlm/configuration_smolvlm.py new file mode 100644 index 0000000000000000000000000000000000000000..067db68e129eb35ecf742eaabe3b0079704c7508 --- /dev/null +++ b/third_party/transformers/src/transformers/models/smolvlm/configuration_smolvlm.py @@ -0,0 +1,118 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/smolvlm/modular_smolvlm.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_smolvlm.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the HuggingFace Inc. team. All rights reserved. +# Written by Orr Zohar +# +# 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, AutoConfig + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="HuggingFaceTB/SmolVLM2-2.2B-Instruct") +@strict +class SmolVLMVisionConfig(PreTrainedConfig): + r""" + Example: + + ```python + >>> from transformers.models.smolvlm.modeling_smolvlm import SmolVLMVisionTransformer + >>> from transformers.models.smolvlm.configuration_smolvlm import SmolVLMVisionConfig + + >>> # Initializing a SmolVLMVisionConfig with google/siglip-so400m-patch14-384 style configuration + >>> configuration = SmolVLMVisionConfig() + + >>> # Initializing a SmolVLMVisionTransformer (with random weights) from the google/siglip-so400m-patch14-384 style configuration + >>> model = SmolVLMVisionTransformer(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "smolvlm_vision" + base_config_key = "vision_config" + + hidden_size: int = 1152 + intermediate_size: int = 3072 + num_hidden_layers: int = 12 + num_attention_heads: int = 16 + 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 = "gelu_pytorch_tanh" + layer_norm_eps: float = 1e-6 + attention_dropout: float | int = 0.0 + initializer_range: float = 0.02 + + +@auto_docstring(checkpoint="HuggingFaceTB/SmolVLM2-2.2B-Instruct") +@strict +class SmolVLMConfig(PreTrainedConfig): + r""" + scale_factor (`int`, *optional*, defaults to 2): + The scale factor for the image encoder. + + Example: + ```python + >>> from transformers import SmolVLMModel, SmolVLMConfig + >>> # Initializing configuration + >>> configuration = SmolVLMConfig() + >>> # Initializing a model from the configuration + >>> model = SmolVLMModel(configuration) + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "smolvlm" + sub_configs = {"text_config": AutoConfig, "vision_config": SmolVLMVisionConfig} + + use_cache: bool = True + image_token_id: int = 128257 + tie_word_embeddings: bool = False + vision_config: dict | PreTrainedConfig | None = None + text_config: dict | PreTrainedConfig | None = None + scale_factor: int = 2 + pad_token_id: int | None = 128_002 + + def __post_init__(self, **kwargs): + if self.vision_config is None: + self.vision_config = SmolVLMVisionConfig() + logger.info("vision_config is None, using default vision config") + elif isinstance(self.vision_config, dict): + self.vision_config = SmolVLMVisionConfig(**self.vision_config) + + if isinstance(self.text_config, dict): + self.text_config["model_type"] = self.text_config.get("model_type", "llama") + self.text_config = CONFIG_MAPPING[self.text_config["model_type"]](**self.text_config) + elif self.text_config is None: + logger.info("text_config is None, using default Llama text config") + self.text_config = CONFIG_MAPPING["llama"]( + rms_norm_eps=1e-5, + pad_token_id=self.pad_token_id, + ) + + super().__post_init__(**kwargs) + + +__all__ = ["SmolVLMVisionConfig", "SmolVLMConfig"] diff --git a/third_party/transformers/src/transformers/models/smolvlm/image_processing_pil_smolvlm.py b/third_party/transformers/src/transformers/models/smolvlm/image_processing_pil_smolvlm.py new file mode 100644 index 0000000000000000000000000000000000000000..3d53ed09c11fd66a4cd1de2d4bd62d7e5ce62f7d --- /dev/null +++ b/third_party/transformers/src/transformers/models/smolvlm/image_processing_pil_smolvlm.py @@ -0,0 +1,481 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/smolvlm/modular_smolvlm.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_smolvlm.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the HuggingFace Inc. team. All rights reserved. +# Written by Orr Zohar +# +# 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 numpy as np + +from ...image_processing_backends import PilBackend +from ...image_processing_utils import BatchFeature +from ...image_transforms import PaddingMode, pad +from ...image_utils import ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ImageInput, + PILImageResampling, + SizeDict, + make_nested_list_of_images, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring + + +class SmolVLMImageProcessorKwargs(ImagesKwargs, total=False): + """ + do_image_splitting (`bool`, *optional*, defaults to `True`): + Whether to split the image into sub-images concatenated with the original image. They are split into patches + such that each patch has a size of `max_image_size["height"]` x `max_image_size["width"]`. + max_image_size (`Dict`, *optional*, defaults to `{"longest_edge": 364}`): + Maximum resolution of the patches of images accepted by the model. This is a dictionary containing the key "longest_edge". + return_row_col_info (`bool`, *optional*, defaults to `False`): + Whether to return the row and column information of the images. + """ + + do_image_splitting: bool + max_image_size: dict[str, int] + return_row_col_info: bool + + +def _make_pixel_mask(image: np.ndarray, output_size: tuple[int, int]) -> np.ndarray: + """Make pixel mask: 1=valid, 0=padding. Images are CHW.""" + h, w = image.shape[-2:] + mask = np.zeros(output_size, dtype=np.int64) + mask[:h, :w] = 1 + return mask + + +# Adapted from transformers.models.smolvlm.image_processing_smolvlm.MAX_IMAGE_SIZE +MAX_IMAGE_SIZE = 4096 # 4k resolution as absolute maximum + + +# Adapted from transformers.models.smolvlm.image_processing_smolvlm._resize_output_size_rescale_to_max_len +def _resize_output_size_rescale_to_max_len( + height: int, width: int, min_len: int | None = 1, max_len: int | None = None +) -> tuple[int, int]: + """ + Get the output size of the image after resizing given a dictionary specifying the max and min sizes. + Args: + height (`int`): + Height of the input image. + width (`int`): + Width of the input image. + min_len (`int`, *optional*, defaults to 1): + Minimum size of the output image. + max_len (`int`, *optional*, defaults to the maximum size of the image): + Maximum size of the output image. + Returns: + The output size of the image after resizing. + """ + max_len = max(height, width) if max_len is None else max_len + aspect_ratio = width / height + + if width >= height: + width = max_len + height = int(width / aspect_ratio) + if height % 2 != 0: + height += 1 + elif height > width: + height = max_len + width = int(height * aspect_ratio) + if width % 2 != 0: + width += 1 + + # Avoid resizing to a size smaller than min_len + height = max(height, min_len) + width = max(width, min_len) + return height, width + + +# Adapted from transformers.models.smolvlm.image_processing_smolvlm._resize_output_size_scale_below_upper_bound +def _resize_output_size_scale_below_upper_bound( + height: int, width: int, max_len: dict[str, int] | None = None +) -> tuple[int, int]: + """ + Get the output size of the image after resizing given a dictionary specifying the max and min sizes. + Args: + height (`int`): + Height of the input image. + width (`int`): + Width of the input image. + max_len (`Dict[str, int]`, *optional*, defaults to the maximum size of the image): + Defines the maximum dimensions of the image. + Returns: + The output size of the image after resizing. + """ + max_len = max(height, width) if max_len is None else max_len + + aspect_ratio = width / height + if width >= height and width > max_len: + width = max_len + height = int(width / aspect_ratio) + elif height > width and height > max_len: + height = max_len + width = int(height * aspect_ratio) + + # Avoid resizing to a size smaller than 1 + height = max(height, 1) + width = max(width, 1) + return height, width + + +def get_max_height_width(images_list: list[list[np.ndarray]]) -> tuple[int, int]: + """ + Get the maximum height and width across all images in a batch. + """ + image_sizes = [] + for images in images_list: + for image in images: + image_sizes.append(image.shape[-2:]) + + max_height = max(size[0] for size in image_sizes) + max_width = max(size[1] for size in image_sizes) + return (max_height, max_width) + + +def get_num_channels(images_list: list[list[np.ndarray]]) -> int: + """ + Get the number of channels across all images in a batch. Handle empty sublists like in [[], [image]]. + """ + for images in images_list: + if images: + return images[0].shape[0] + + raise ValueError("No images found in the batch.") + + +def get_resize_output_image_size( + image: np.ndarray, + resolution_max_side: int, +) -> tuple[int, int]: + """ + Get the output size of the image after resizing given a dictionary specifying the max and min sizes. + Args: + image (`np.ndarray`): + Image to resize. + resolution_max_side (`int`): + The longest edge of the image will be resized to this value. The shortest edge will be resized to keep the + input aspect ratio. + Returns: + The output size of the image after resizing. + """ + height, width = image.shape[-2:] + + # Find the output size, when rescaling the longest edge to max_len and preserving the aspect ratio + height, width = _resize_output_size_rescale_to_max_len(height, width, max_len=resolution_max_side) + # Find the output size when scaling the image to be below the MAX_IMAGE_SIZE + height, width = _resize_output_size_scale_below_upper_bound(height, width, max_len=MAX_IMAGE_SIZE) + return height, width + + +@auto_docstring +class SmolVLMImageProcessorPil(PilBackend): + resample = PILImageResampling.LANCZOS + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"longest_edge": 4 * 364} + max_image_size = {"longest_edge": 364} + do_resize = True + do_rescale = True + do_normalize = True + do_convert_rgb = True + do_image_splitting = True + do_pad = True + return_row_col_info = False + valid_kwargs = SmolVLMImageProcessorKwargs + model_input_names = ["pixel_values", "pixel_attention_mask"] + + def __init__(self, **kwargs: Unpack[SmolVLMImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[SmolVLMImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def _prepare_images_structure(self, images: ImageInput, expected_ndims: int = 3) -> ImageInput: + images = self.fetch_images(images) + return make_nested_list_of_images(images, expected_ndims=expected_ndims) + + def resize( + self, + image: np.ndarray, + size: SizeDict, + resample: PILImageResampling = PILImageResampling.LANCZOS, + **kwargs, + ) -> np.ndarray: + if size.longest_edge: + new_size = get_resize_output_image_size(image, resolution_max_side=size.longest_edge) + elif size.height and size.width: + new_size = (size.height, size.width) + else: + raise ValueError("size must be a dictionary with key 'longest_edge' or 'height' and 'width'.") + return super().resize(image, SizeDict(height=new_size[0], width=new_size[1]), resample=resample, **kwargs) + + def split_images( + self, + image: np.ndarray, + max_image_size: dict[str, int], + resample: "PILImageResampling | None" = None, + ): + """Split an image into patches (mirrors TorchvisionBackend.split_images). Images are always CHW.""" + num_channels, height, width = image.shape + max_height = max_width = max_image_size["longest_edge"] + + if height > max_height or width > max_width: + num_splits_h = (height - max_height) // max_height + 1 + num_splits_w = (width - max_width) // max_width + 1 + + frames = [] + for r in range(num_splits_h): + for c in range(num_splits_w): + start_y = r * max_height + start_x = c * max_width + end_y = start_y + max_height + end_x = start_x + max_width + crop = image[:, start_y:end_y, start_x:end_x] + frames.append(crop) + + global_image_height, global_image_width = max_height, max_width + image = self.resize( + image, SizeDict(height=global_image_height, width=global_image_width), resample=resample + ) + frames.append(image) + else: + num_splits_h, num_splits_w = 0, 0 + frames = [image] + + return frames, num_splits_h, num_splits_w + + def resize_for_vision_encoder( + self, + image: np.ndarray, + vision_encoder_max_size: int, + resample: "PILImageResampling | None" = None, + ): + """Resize images to be multiples of vision_encoder_max_size. Images are always CHW.""" + height, width = image.shape[-2:] + aspect_ratio = width / height + if width >= height: + width = math.ceil(width / vision_encoder_max_size) * vision_encoder_max_size + height = int(width / aspect_ratio) + height = math.ceil(height / vision_encoder_max_size) * vision_encoder_max_size + elif height > width: + height = math.ceil(height / vision_encoder_max_size) * vision_encoder_max_size + width = int(height * aspect_ratio) + width = math.ceil(width / vision_encoder_max_size) * vision_encoder_max_size + new_size = SizeDict(height=height, width=width) + return self.resize(image, size=new_size, resample=resample) + + def pad( + self, + image: np.ndarray, + padded_size: tuple[int, int], + fill: int = 0, + return_pixel_mask: bool = True, + ): + """Pad image to padded_size. Mirrors TorchvisionBackend.pad. Images are always CHW.""" + original_size = image.shape[-2:] + padding_bottom = padded_size[0] - original_size[0] + padding_right = padded_size[1] - original_size[1] + + if padding_bottom < 0 or padding_right < 0: + raise ValueError( + f"Padding dimensions are negative. Please make sure that the padded size is larger than the " + f"original size. Got padded size: {padded_size}, original size: {original_size}." + ) + + pixel_mask = _make_pixel_mask(image, output_size=padded_size) if return_pixel_mask else None + + if original_size != padded_size: + padding = ((0, padding_bottom), (0, padding_right)) + image = pad( + image, + padding, + mode=PaddingMode.CONSTANT, + constant_values=fill, + data_format="channels_first", + input_data_format="channels_first", + ) + + return image, pixel_mask + + def _preprocess( + self, + images: list[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, + do_pad: bool | None, + do_image_splitting: bool | None, + max_image_size: dict[str, int] | None, + return_row_col_info: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + """Process a batch of images. Mirrors TorchvisionBackend._preprocess with per-image loops instead of batching.""" + # Resize + if do_resize: + images = [ + [self.resize(image=img, size=size, resample=resample) for img in batch_images] + for batch_images in images + ] + + # Image splitting + if do_image_splitting: + images = [ + [ + self.resize_for_vision_encoder(image, max_image_size["longest_edge"], resample=resample) + for image in batch_images + ] + for batch_images in images + ] + images_split_arrays = [] + images_rows = [] + images_cols = [] + for batch_images in images: + split_image_arrays = [] + image_rows = [] + image_cols = [] + for image in batch_images: + split_image_array, rows, cols = self.split_images( + image, max_image_size=max_image_size, resample=resample + ) + split_image_arrays.extend(split_image_array) + image_rows.append(rows) + image_cols.append(cols) + images_split_arrays.append(split_image_arrays) + images_rows.append(image_rows) + images_cols.append(image_cols) + images = images_split_arrays + rows = images_rows + cols = images_cols + else: + images = [ + [ + self.resize( + image=image, + size=SizeDict(height=max_image_size["longest_edge"], width=max_image_size["longest_edge"]), + resample=resample, + ) + for image in batch_images + ] + for batch_images in images + ] + rows = [[0] * len(batch_images) for batch_images in images] + cols = [[0] * len(batch_images) for batch_images in images] + + # Rescale and normalize + if do_rescale: + images = [[self.rescale(img, rescale_factor) for img in batch_images] for batch_images in images] + if do_normalize: + images = [[self.normalize(img, image_mean, image_std) for img in batch_images] for batch_images in images] + + # Pad + if do_pad: + max_num_images = max(len(images_) for images_ in images) + max_height, max_width = get_max_height_width(images) + num_channels = get_num_channels(images) + + padded_images_list = [ + [np.zeros((num_channels, max_height, max_width), dtype=np.float32) for _ in range(max_num_images)] + for _ in range(len(images)) + ] + pixel_attention_masks = [ + [np.zeros((max_height, max_width), dtype=np.int64) for _ in range(max_num_images)] + for _ in range(len(images)) + ] + + for i, batch_images in enumerate(images): + for j, image in enumerate(batch_images): + padded_images_list[i][j], pixel_attention_masks[i][j] = self.pad(image, (max_height, max_width)) + images = padded_images_list + + if do_pad: + data = { + "pixel_values": np.array(images), + "pixel_attention_mask": np.array(pixel_attention_masks), + } + elif return_tensors == "pt": + data = {"pixel_values": np.asarray(images)} + else: + data = {"pixel_values": images} + + encoding = BatchFeature(data=data, tensor_type=return_tensors) + if return_row_col_info: + encoding["rows"] = rows + encoding["cols"] = cols + + return encoding + + def to_dict(self): + encoder_dict = super().to_dict() + encoder_dict.pop("_valid_processor_keys", None) + encoder_dict.pop("return_row_col_info", None) + return encoder_dict + + def get_number_of_image_patches(self, height: int, width: int, images_kwargs: dict): + """ + A utility that returns number of image patches for a given image size. + + Args: + height (`int`): + Height of the input image. + width (`int`): + Width of the input image. + images_kwargs (`dict`) + Any kwargs to override defaults of the image processor. + Returns: + `int`: Number of patches per image. + """ + do_image_splitting = images_kwargs.get("do_image_splitting", self.do_image_splitting) + max_image_size = images_kwargs.get("max_image_size", self.max_image_size) + size = images_kwargs.get("size", self.size) + + num_patches = num_rows = num_cols = 0 + if do_image_splitting: + height, width = _resize_output_size_rescale_to_max_len(height, width, max_len=size["longest_edge"]) + height, width = _resize_output_size_scale_below_upper_bound(height, width, max_len=MAX_IMAGE_SIZE) + aspect_ratio = width / height + + if width >= height: + resized_width = math.ceil(width / max_image_size["longest_edge"]) * max_image_size["longest_edge"] + resized_height = int(width / aspect_ratio) + resized_height = math.ceil(height / max_image_size["longest_edge"]) * max_image_size["longest_edge"] + elif height > width: + resized_height = math.ceil(height / max_image_size["longest_edge"]) * max_image_size["longest_edge"] + resized_width = int(height * aspect_ratio) + resized_width = math.ceil(width / max_image_size["longest_edge"]) * max_image_size["longest_edge"] + + max_height = max_width = max_image_size["longest_edge"] + if resized_height > max_height or resized_width > max_width: + num_rows = math.ceil(resized_height / max_height) + num_cols = math.ceil(resized_width / max_width) + num_patches = num_rows * num_cols + 1 + + return num_patches, num_rows, num_cols + + +__all__ = ["SmolVLMImageProcessorPil"] diff --git a/third_party/transformers/src/transformers/models/smolvlm/image_processing_smolvlm.py b/third_party/transformers/src/transformers/models/smolvlm/image_processing_smolvlm.py new file mode 100644 index 0000000000000000000000000000000000000000..6c22e496c6697387931240a5c6f2f1222a73cd41 --- /dev/null +++ b/third_party/transformers/src/transformers/models/smolvlm/image_processing_smolvlm.py @@ -0,0 +1,540 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/smolvlm/modular_smolvlm.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_smolvlm.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the HuggingFace Inc. team. All rights reserved. +# Written by Orr Zohar +# +# 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 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 ( + IMAGENET_STANDARD_MEAN, + IMAGENET_STANDARD_STD, + ImageInput, + PILImageResampling, + SizeDict, + make_nested_list_of_images, +) +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TensorType, auto_docstring + + +class SmolVLMImageProcessorKwargs(ImagesKwargs, total=False): + """ + do_image_splitting (`bool`, *optional*, defaults to `True`): + Whether to split the image into sub-images concatenated with the original image. They are split into patches + such that each patch has a size of `max_image_size["height"]` x `max_image_size["width"]`. + max_image_size (`Dict`, *optional*, defaults to `{"longest_edge": 364}`): + Maximum resolution of the patches of images accepted by the model. This is a dictionary containing the key "longest_edge". + return_row_col_info (`bool`, *optional*, defaults to `False`): + Whether to return the row and column information of the images. + """ + + do_image_splitting: bool + max_image_size: dict[str, int] + return_row_col_info: bool + + +MAX_IMAGE_SIZE = 4096 # 4k resolution as absolute maximum + + +def _resize_output_size_rescale_to_max_len( + height: int, width: int, min_len: int | None = 1, max_len: int | None = None +) -> tuple[int, int]: + """ + Get the output size of the image after resizing given a dictionary specifying the max and min sizes. + Args: + height (`int`): + Height of the input image. + width (`int`): + Width of the input image. + min_len (`int`, *optional*, defaults to 1): + Minimum size of the output image. + max_len (`int`, *optional*, defaults to the maximum size of the image): + Maximum size of the output image. + Returns: + The output size of the image after resizing. + """ + max_len = max(height, width) if max_len is None else max_len + aspect_ratio = width / height + + if width >= height: + width = max_len + height = int(width / aspect_ratio) + if height % 2 != 0: + height += 1 + elif height > width: + height = max_len + width = int(height * aspect_ratio) + if width % 2 != 0: + width += 1 + + # Avoid resizing to a size smaller than min_len + height = max(height, min_len) + width = max(width, min_len) + return height, width + + +def _resize_output_size_scale_below_upper_bound( + height: int, width: int, max_len: dict[str, int] | None = None +) -> tuple[int, int]: + """ + Get the output size of the image after resizing given a dictionary specifying the max and min sizes. + Args: + height (`int`): + Height of the input image. + width (`int`): + Width of the input image. + max_len (`Dict[str, int]`, *optional*, defaults to the maximum size of the image): + Defines the maximum dimensions of the image. + Returns: + The output size of the image after resizing. + """ + max_len = max(height, width) if max_len is None else max_len + + aspect_ratio = width / height + if width >= height and width > max_len: + width = max_len + height = int(width / aspect_ratio) + elif height > width and height > max_len: + height = max_len + width = int(height * aspect_ratio) + + # Avoid resizing to a size smaller than 1 + height = max(height, 1) + width = max(width, 1) + return height, width + + +def get_resize_output_image_size( + image: "torch.Tensor", + resolution_max_side: int, +) -> tuple[int, int]: + """ + Get the output size of the image after resizing given a dictionary specifying the max and min sizes. + Args: + image (`torch.Tensor`): + Image to resize. + resolution_max_side (`int`): + The longest edge of the image will be resized to this value. The shortest edge will be resized to keep the + input aspect ratio. + Returns: + The output size of the image after resizing. + """ + height, width = image.shape[-2:] + + # Find the output size, when rescaling the longest edge to max_len and preserving the aspect ratio + height, width = _resize_output_size_rescale_to_max_len(height, width, max_len=resolution_max_side) + # Find the output size when scaling the image to be below the MAX_IMAGE_SIZE + height, width = _resize_output_size_scale_below_upper_bound(height, width, max_len=MAX_IMAGE_SIZE) + return height, width + + +def get_max_height_width(images_list: list[list["torch.Tensor|np.ndarray"]]) -> tuple[int, int]: + """ + Get the maximum height and width across all images in a batch. + """ + image_sizes = [] + for images in images_list: + for image in images: + image_sizes.append(image.shape[-2:]) + + max_height = max(size[0] for size in image_sizes) + max_width = max(size[1] for size in image_sizes) + return (max_height, max_width) + + +def get_num_channels(images_list: list[list["torch.Tensor|np.ndarray"]]) -> int: + """ + Get the number of channels across all images in a batch. Handle empty sublists like in [[], [image]]. + """ + for images in images_list: + if images: + return images[0].shape[0] + + raise ValueError("No images found in the batch.") + + +def get_device_from_images(images_list: list[list["torch.Tensor"]]) -> "torch.device": + """ + Get the device from the first non-empty element in a nested list of images. + Handle empty sublists like in [[], [image]]. + """ + for images in images_list: + if images: + return images[0].device + + +@auto_docstring +class SmolVLMImageProcessor(TorchvisionBackend): + resample = PILImageResampling.LANCZOS + image_mean = IMAGENET_STANDARD_MEAN + image_std = IMAGENET_STANDARD_STD + size = {"longest_edge": 4 * 364} + max_image_size = {"longest_edge": 364} + do_resize = True + do_rescale = True + do_normalize = True + do_convert_rgb = True + do_image_splitting = True + do_pad = True + return_row_col_info = False + valid_kwargs = SmolVLMImageProcessorKwargs + model_input_names = ["pixel_values", "pixel_attention_mask"] + + def __init__(self, **kwargs: Unpack[SmolVLMImageProcessorKwargs]): + super().__init__(**kwargs) + + @auto_docstring + def preprocess(self, images: ImageInput, **kwargs: Unpack[SmolVLMImageProcessorKwargs]) -> BatchFeature: + return super().preprocess(images, **kwargs) + + def _prepare_images_structure(self, images: ImageInput, expected_ndims: int = 3) -> ImageInput: + """ + Prepare a nested images structure for processing. + """ + # Checks for `str` in case of URL/local path and optionally loads images + images = self.fetch_images(images) + return make_nested_list_of_images(images, expected_ndims=expected_ndims) + + def resize( + self, + image: "torch.Tensor", + size: SizeDict, + resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None, + **kwargs, + ) -> "torch.Tensor": + """ + Resize an image. The longest edge of the image is resized to size.longest_edge, with the shortest edge + resized to keep the input aspect ratio. Can also be used with size.height and size.width. + Args: + image (`torch.Tensor`): + Image to resize. + size (`SizeDict`): + Size of the output image. + resample (`PILImageResampling | tvF.InterpolationMode | int | None`, *optional*): + Resampling filter to use when resizing the image. + """ + if size.longest_edge: + new_size = get_resize_output_image_size(image, resolution_max_side=size.longest_edge) + elif size.height and size.width: + new_size = (size.height, size.width) + else: + raise ValueError("size must be a dictionary with key 'longest_edge' or 'height' and 'width'.") + + return super().resize(image, SizeDict(height=new_size[0], width=new_size[1]), resample=resample, **kwargs) + + def split_images( + self, + images: torch.Tensor, + max_image_size: dict[str, int], + resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None, + ): + """ + Split an image into squares of side max_image_size and the original image resized to max_image_size. + That means that a single image becomes a sequence of images. + This is a "trick" to spend more compute on each image with no changes in the vision encoder. + 1) If one side of the original image is larger than `max_image_size`, resize it to `max_image_size` while preserving the aspect ratio. + 2) Divide the resulting image into `ceil(height / max_image_size)` x `ceil(width / max_image_size)` + sub-images of the same size each (image_size, image_size). Typically, 364x364. + 3) Returns the list of the crops and the original image, in addition to the number of splits for the height and the width. + Args: + images (`torch.Tensor`): + Images to split. + max_image_size (`Dict[str, int]`): + Maximum size of the output image. If the image is larger than this size, it will be split into + patches of this size, and the original image will be concatenated with the patches, resized to max_size. + resample (`PILImageResampling | tvF.InterpolationMode | int | None`, *optional*): + Resampling filter to use when resizing the image. + """ + batch_size, num_channels, height, width = images.size() + height_dim, width_dim = 2, 3 + + max_height = max_width = max_image_size["longest_edge"] + + frames = [] + if height > max_height or width > max_width: + # Calculate the number of splits + num_splits_h = math.ceil(height / max_height) + num_splits_w = math.ceil(width / max_width) + + # Split the images by height, then by width + frames = ( + images.unfold(height_dim, size=max_height, step=max_height) + .unfold(width_dim, size=max_width, step=max_width) + .contiguous() + .view(batch_size, num_channels, -1, max_height, max_width) + .permute(0, 2, 1, 3, 4) + ) # batch_size x n_frames x num_channels x height x width + + # For the global image at the end, we resize it to match the max_image_size, for cpu memory efficiency + global_image_height, global_image_width = max_height, max_width + images = self.resize( + images, SizeDict(height=global_image_height, width=global_image_width), resample=resample + ) + + frames = torch.cat((frames, images.unsqueeze(1)), dim=1) + else: + num_splits_h, num_splits_w = 0, 0 + frames = images.unsqueeze(1) + + num_splits_h = [num_splits_h] * batch_size + num_splits_w = [num_splits_w] * batch_size + + return frames, num_splits_h, num_splits_w + + def resize_for_vision_encoder( + self, + image: torch.Tensor, + vision_encoder_max_size: int, + resample: "PILImageResampling | tvF.InterpolationMode | int | None" = None, + ): + """ + Resize images to be multiples of `vision_encoder_max_size` while preserving the aspect ratio. + Args: + image (`torch.Tensor`): + Images to resize. + vision_encoder_max_size (`int`): + Maximum size of the output image. If the image is larger than this size, it will be split into + patches of this size, and the original image will be concatenated with the patches, resized to max_size. + resample (`PILImageResampling | tvF.InterpolationMode | int | None`, *optional*): + Resampling filter to use when resizing the image. + """ + height, width = image.size()[-2:] + + aspect_ratio = width / height + if width >= height: + width = math.ceil(width / vision_encoder_max_size) * vision_encoder_max_size + height = int(width / aspect_ratio) + height = math.ceil(height / vision_encoder_max_size) * vision_encoder_max_size + elif height > width: + height = math.ceil(height / vision_encoder_max_size) * vision_encoder_max_size + width = int(height * aspect_ratio) + width = math.ceil(width / vision_encoder_max_size) * vision_encoder_max_size + new_size = SizeDict(height=height, width=width) + return self.resize(image, size=new_size, resample=resample) + + def pad( + self, + image: torch.Tensor, + padded_size: tuple[int, int], + fill: int = 0, + return_pixel_mask: bool = True, + ): + original_size = image.shape[-2:] + padding_bottom = padded_size[0] - original_size[0] + padding_right = padded_size[1] - original_size[1] + + if padding_bottom < 0 or padding_right < 0: + raise ValueError( + f"Padding dimensions are negative. Please make sure that the padded size is larger than the " + f"original size. Got padded size: {padded_size}, original size: {original_size}." + ) + + # Only pad if necessary + if original_size != padded_size: + padding = (0, 0, padding_right, padding_bottom) + image = tvF.pad(image, padding, fill=fill, padding_mode="constant") + + # Make a pixel mask for the image, where 1 indicates a valid pixel and 0 indicates padding. + pixel_mask = None + if return_pixel_mask: + pixel_mask = torch.zeros_like(image[..., 0, :, :], dtype=torch.int64) + pixel_mask[: original_size[0], : original_size[1]] = 1 + + return image, pixel_mask + + def _preprocess( + self, + images: list[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, + do_pad: bool | None, + do_image_splitting: bool | None, + max_image_size: dict[str, int] | None, + return_row_col_info: bool | None, + disable_grouping: bool | None, + return_tensors: str | TensorType | None, + **kwargs, + ) -> BatchFeature: + """ + Process a batch of images for the model. + """ + + grouped_images, grouped_images_index = group_images_by_shape( + images, is_nested=True, 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=resample) + resized_images_grouped[shape] = stacked_images + resized_images = reorder_images(resized_images_grouped, grouped_images_index, is_nested=True) + + grouped_images, grouped_images_index = group_images_by_shape( + resized_images, is_nested=True, disable_grouping=disable_grouping + ) + split_images_grouped = {} + if do_image_splitting: + rows_grouped = {} + cols_grouped = {} + for shape, stacked_images in grouped_images.items(): + stacked_images = self.resize_for_vision_encoder( + stacked_images, max_image_size["longest_edge"], resample=resample + ) + stacked_images, rows, cols = self.split_images( + stacked_images, max_image_size=max_image_size, resample=resample + ) + split_images_grouped[shape] = stacked_images + rows_grouped[shape] = rows + cols_grouped[shape] = cols + processed_images = reorder_images(split_images_grouped, grouped_images_index, is_nested=True) + rows = reorder_images(rows_grouped, grouped_images_index, is_nested=True) + cols = reorder_images(cols_grouped, grouped_images_index, is_nested=True) + # flattenened the doubly nested list to a nested list + for i, group_images in enumerate(processed_images): + processed_images[i] = [image for sublist in group_images for image in sublist] + else: + for shape, stacked_images in grouped_images.items(): + # We square the images to max_image_size + stacked_images = self.resize( + image=stacked_images, + size=SizeDict(height=max_image_size["longest_edge"], width=max_image_size["longest_edge"]), + resample=resample, + ) + split_images_grouped[shape] = stacked_images + processed_images = reorder_images(split_images_grouped, grouped_images_index, is_nested=True) + rows = [[0] * len(images) for images in processed_images] + cols = [[0] * len(images) for images in processed_images] + # 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( + processed_images, is_nested=True, disable_grouping=disable_grouping + ) + processed_images_grouped = {} + for shape, stacked_images in grouped_images.items(): + # 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, is_nested=True) + if do_pad: + # Get max images per batch + max_num_images = max(len(images_) for images_ in processed_images) + max_height, max_width = get_max_height_width(processed_images) + num_channels = get_num_channels(processed_images) + device = get_device_from_images(processed_images) + + processed_images_padded = torch.zeros( + len(processed_images), + max_num_images, + *(num_channels, max_height, max_width), + device=device, + ) + pixel_attention_masks = torch.zeros( + len(processed_images), + max_num_images, + *(max_height, max_width), + device=device, + ) + for i, images in enumerate(processed_images): + for j, image in enumerate(images): + processed_images_padded[i, j], pixel_attention_masks[i, j] = self.pad( + image, (max_height, max_width) + ) + processed_images = processed_images_padded + + if do_pad: + data = {"pixel_values": processed_images, "pixel_attention_mask": pixel_attention_masks} + elif return_tensors == "pt": + data = {"pixel_values": torch.stack([torch.stack(images) for images in processed_images])} + else: + data = {"pixel_values": processed_images} + # This is needed for generating correct text inputs in the processor - we don't pad to the max number of images + encoding = BatchFeature(data=data, tensor_type=return_tensors) + + if return_row_col_info: + encoding["rows"] = rows + encoding["cols"] = cols + + return encoding + + def to_dict(self): + encoder_dict = super().to_dict() + encoder_dict.pop("_valid_processor_keys", None) + encoder_dict.pop("return_row_col_info", None) + return encoder_dict + + def get_number_of_image_patches(self, height: int, width: int, images_kwargs: dict): + """ + A utility that returns number of image patches for a given image size. + + Args: + height (`int`): + Height of the input image. + width (`int`): + Width of the input image. + images_kwargs (`dict`) + Any kwargs to override defaults of the image processor. + Returns: + `int`: Number of patches per image. + """ + do_image_splitting = images_kwargs.get("do_image_splitting", self.do_image_splitting) + max_image_size = images_kwargs.get("max_image_size", self.max_image_size) + size = images_kwargs.get("size", self.size) + + num_patches = num_rows = num_cols = 0 + if do_image_splitting: + height, width = _resize_output_size_rescale_to_max_len(height, width, max_len=size["longest_edge"]) + height, width = _resize_output_size_scale_below_upper_bound(height, width, max_len=MAX_IMAGE_SIZE) + aspect_ratio = width / height + + if width >= height: + resized_width = math.ceil(width / max_image_size["longest_edge"]) * max_image_size["longest_edge"] + resized_height = int(width / aspect_ratio) + resized_height = math.ceil(height / max_image_size["longest_edge"]) * max_image_size["longest_edge"] + elif height > width: + resized_height = math.ceil(height / max_image_size["longest_edge"]) * max_image_size["longest_edge"] + resized_width = int(height * aspect_ratio) + resized_width = math.ceil(width / max_image_size["longest_edge"]) * max_image_size["longest_edge"] + + max_height = max_width = max_image_size["longest_edge"] + if resized_height > max_height or resized_width > max_width: + # Calculate the number of splits + num_rows = math.ceil(resized_height / max_height) + num_cols = math.ceil(resized_width / max_width) + num_patches = num_rows * num_cols + 1 + + return num_patches, num_rows, num_cols + + +__all__ = ["SmolVLMImageProcessor"] diff --git a/third_party/transformers/src/transformers/models/smolvlm/modeling_smolvlm.py b/third_party/transformers/src/transformers/models/smolvlm/modeling_smolvlm.py new file mode 100644 index 0000000000000000000000000000000000000000..ed2eec345b75437d3f7c58b1503ed71a68c2f026 --- /dev/null +++ b/third_party/transformers/src/transformers/models/smolvlm/modeling_smolvlm.py @@ -0,0 +1,876 @@ +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# This file was automatically generated from src/transformers/models/smolvlm/modular_smolvlm.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_smolvlm.py file directly. One of our CI enforces this. +# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨 +# Copyright 2025 the HuggingFace Inc. team. All rights reserved. +# Written by Orr Zohar +# +# 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 + +import torch +from torch import nn + +from ...activations import ACT2FN +from ...cache_utils import Cache, DynamicCache +from ...generation import GenerationConfig, GenerationMixin +from ...masking_utils import create_bidirectional_mask +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_layers import GradientCheckpointingLayer +from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, ModelOutput +from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel +from ...processing_utils import Unpack +from ...utils import TransformersKwargs, 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 ..auto import AutoModel +from .configuration_smolvlm import SmolVLMConfig, SmolVLMVisionConfig + + +logger = logging.get_logger(__name__) + + +@auto_docstring +class SmolVLMPreTrainedModel(PreTrainedModel): + config: SmolVLMConfig + base_model_prefix = "model" + input_modalities = ("image", "text") + supports_gradient_checkpointing = True + _no_split_modules = ["SmolVLMVisionAttention", "SmolVLMDecoderLayer"] + _skip_keys_device_placement = "past_key_values" + _supports_flash_attn = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_attention_backend = True + + +class SmolVLMVisionEmbeddings(nn.Module): + """ + This is a modified version of `siglip.modelign_siglip.SiglipVisionEmbeddings` to enable images of variable + resolution. + + The modifications are adapted from [Patch n' Pack: NaViT, a Vision Transformer for any Aspect Ratio and Resolution](https://huggingface.co/papers/2307.06304) + which allows treating images in their native aspect ratio and without the need to resize them to the same + fixed size. In particular, we start from the original pre-trained SigLIP model + (which uses images of fixed-size square images) and adapt it by training on images of variable resolutions. + """ + + def __init__(self, config: SmolVLMVisionConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.image_size = config.image_size + self.patch_size = config.patch_size + + self.patch_embedding = nn.Conv2d( + in_channels=config.num_channels, + out_channels=self.embed_dim, + kernel_size=self.patch_size, + stride=self.patch_size, + padding="valid", + ) + + self.num_patches_per_side = self.image_size // self.patch_size + self.num_patches = self.num_patches_per_side**2 + self.num_positions = self.num_patches + self.position_embedding = nn.Embedding(self.num_positions, self.embed_dim) + + def forward(self, pixel_values: torch.FloatTensor, patch_attention_mask: torch.BoolTensor) -> torch.Tensor: + batch_size, _, max_im_h, max_im_w = pixel_values.shape + + patch_embeds = self.patch_embedding(pixel_values) + embeddings = patch_embeds.flatten(2).transpose(1, 2) + + max_nb_patches_h, max_nb_patches_w = max_im_h // self.patch_size, max_im_w // self.patch_size + boundaries = torch.arange( + 1 / self.num_patches_per_side, 1.0, 1 / self.num_patches_per_side, device=pixel_values.device + ) + position_ids = torch.full( + size=(batch_size, max_nb_patches_h * max_nb_patches_w), fill_value=0, device=pixel_values.device + ) + + nb_patches_h = patch_attention_mask[:, :, 0].sum(dim=1) # (batch_size,) + nb_patches_w = patch_attention_mask[:, 0, :].sum(dim=1) # (batch_size,) + + step_h = 1.0 / nb_patches_h # (batch_size,) + step_w = 1.0 / nb_patches_w # (batch_size,) + + max_patches_h = patch_attention_mask.size(1) + max_patches_w = patch_attention_mask.size(2) + h_indices = torch.arange(max_patches_h, device=position_ids.device, dtype=torch.float32) + w_indices = torch.arange(max_patches_w, device=position_ids.device, dtype=torch.float32) + + fractional_coords_h = h_indices[None, :] * step_h[:, None] + fractional_coords_w = w_indices[None, :] * step_w[:, None] + + fractional_coords_h = torch.clamp(fractional_coords_h, max=(1.0 - 1e-6)) + fractional_coords_w = torch.clamp(fractional_coords_w, max=(1.0 - 1e-6)) + + fractional_coords_h = fractional_coords_h.to(pixel_values.dtype) + fractional_coords_w = fractional_coords_w.to(pixel_values.dtype) + + bucket_coords_h = torch.bucketize(fractional_coords_h, boundaries, right=True) + bucket_coords_w = torch.bucketize(fractional_coords_w, boundaries, right=True) + + pos_ids = bucket_coords_h[:, :, None] * self.num_patches_per_side + bucket_coords_w[:, None, :] + pos_ids = pos_ids.reshape(batch_size, -1) + + position_ids[patch_attention_mask.view(batch_size, -1)] = pos_ids[patch_attention_mask.view(batch_size, -1)] + + embeddings = embeddings + self.position_embedding(position_ids) + 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 SmolVLMVisionAttention(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) + + # Ignore copy + self.is_causal = False + + 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 SmolVLMVisionMLP(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 SmolVLMEncoderLayer(GradientCheckpointingLayer): + def __init__(self, config: SmolVLMVisionConfig): + super().__init__() + self.embed_dim = config.hidden_size + self.self_attn = SmolVLMVisionAttention(config) + self.layer_norm1 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + self.mlp = SmolVLMVisionMLP(config) + self.layer_norm2 = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_eps) + + @auto_docstring + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: 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, + 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 SmolVLMEncoder(nn.Module): + """ + Transformer encoder consisting of `config.num_hidden_layers` self attention layers. Each layer is a + [`SmolVLMEncoderLayer`]. + + Args: + config: SmolVLMConfig + """ + + def __init__(self, config: SmolVLMConfig): + super().__init__() + self.config = config + self.layers = nn.ModuleList([SmolVLMEncoderLayer(config) for _ in range(config.num_hidden_layers)]) + self.gradient_checkpointing = False + + # Ignore copy + @auto_docstring + def forward( + self, + inputs_embeds, + attention_mask: torch.Tensor | None = None, + ) -> tuple | BaseModelOutput: + hidden_states = inputs_embeds + for encoder_layer in self.layers: + layer_outputs = encoder_layer( + hidden_states, + attention_mask, + ) + + hidden_states = layer_outputs + + return BaseModelOutput(last_hidden_state=hidden_states) + + +@auto_docstring( + custom_intro=""" + The SmolVLM Vision Transformer Model outputting raw image embedding. + """ +) +class SmolVLMVisionTransformer(SmolVLMPreTrainedModel): + config: SmolVLMVisionConfig + input_modalities = ("image",) + _can_record_outputs = { + "hidden_states": SmolVLMEncoderLayer, + "attentions": SmolVLMVisionAttention, + } + + def __init__(self, config: SmolVLMVisionConfig): + super().__init__(config) + embed_dim = config.hidden_size + + self.embeddings = SmolVLMVisionEmbeddings(config) + self.encoder = SmolVLMEncoder(config) + self.patch_size = config.patch_size + self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps) + + self.post_init() + + def get_input_embeddings(self): + return self.embeddings + + def set_input_embeddings(self, value): + self.embeddings = value + + @merge_with_config_defaults + @capture_outputs(tie_last_hidden_states=False) + def forward( + self, + pixel_values, + patch_attention_mask: torch.BoolTensor | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> tuple | BaseModelOutput: + batch_size = pixel_values.size(0) + if patch_attention_mask is None: + patch_size = self.patch_size + patch_attention_mask = torch.ones( + ( + batch_size, + pixel_values.size(2) // patch_size, + pixel_values.size(3) // patch_size, + ) + ) + patch_attention_mask = patch_attention_mask.to(dtype=torch.bool, device=pixel_values.device) + + hidden_states = self.embeddings(pixel_values=pixel_values, patch_attention_mask=patch_attention_mask) + + patch_attention_mask = patch_attention_mask.view(batch_size, -1) + # Create the correct attention mask based on the attention implementation + patch_attention_mask = create_bidirectional_mask( + config=self.config, + inputs_embeds=hidden_states, + attention_mask=patch_attention_mask, + ) + + encoder_outputs: BaseModelOutput = self.encoder( + inputs_embeds=hidden_states, + attention_mask=patch_attention_mask, + ) + + last_hidden_state = encoder_outputs.last_hidden_state + last_hidden_state = self.post_layernorm(last_hidden_state) + + return BaseModelOutput( + last_hidden_state=last_hidden_state, + ) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for SmolVLM model's outputs that may also contain a past key/values (to speed up sequential decoding). + """ +) +class SmolVLMBaseModelOutputWithPast(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 + """ + + 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 + + +class SmolVLMSimpleMLP(nn.Module): + def __init__(self, config): + super().__init__() + input_size = config.vision_config.hidden_size * (config.scale_factor**2) + output_size = config.text_config.hidden_size + self.proj = nn.Linear(input_size, output_size, bias=False) + + def forward(self, x): + return self.proj(x) + + +class SmolVLMConnector(nn.Module): + def __init__(self, config): + super().__init__() + self.scale_factor = config.scale_factor + self.modality_projection = SmolVLMSimpleMLP(config) + + def pixel_shuffle(self, x, scale_factor=2): + bsz, seq, embed_dim = x.size() + height = width = int(seq**0.5) + x = x.view(bsz, height, width, embed_dim) + x = x.view(bsz, height, int(width / scale_factor), embed_dim * scale_factor) + x = x.permute(0, 2, 1, 3) + x = x.reshape(bsz, int(width / scale_factor), int(height / scale_factor), embed_dim * (scale_factor**2)) + x = x.permute(0, 2, 1, 3) + x = x.reshape(bsz, int(seq / (scale_factor**2)), embed_dim * (scale_factor**2)) + return x + + def forward(self, image_hidden_states): + image_hidden_states = self.pixel_shuffle(image_hidden_states, self.scale_factor) + image_hidden_states = self.modality_projection(image_hidden_states) + return image_hidden_states + + +@auto_docstring( + custom_intro=""" + SmolVLM model consisting of a SIGLIP vision encoder and Llama3 language decoder + """ +) +class SmolVLMModel(SmolVLMPreTrainedModel): + """ + A subclass of Idefics3Model. We do *not* remove or block the call to inputs_merger + in forward. Instead, we override inputs_merger here with custom logic. + """ + + def __init__(self, config: SmolVLMConfig): + super().__init__(config) + self.padding_idx = self.config.text_config.pad_token_id + self.vocab_size = self.config.text_config.vocab_size + + self.vision_model = SmolVLMVisionTransformer._from_config(config.vision_config) + self.connector = SmolVLMConnector(config) + self.text_model = AutoModel.from_config(config.text_config) + + self.image_seq_len = int( + ((config.vision_config.image_size // config.vision_config.patch_size) ** 2) / (config.scale_factor**2) + ) + self.image_token_id = self.config.image_token_id + + 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 inputs_merger( + self, input_ids: torch.LongTensor, inputs_embeds: torch.Tensor, image_hidden_states: torch.Tensor + ): + """ + This method aims at merging the token embeddings with the image hidden states into one single sequence of vectors that are fed to the transformer LM. + The merging happens as follows: + - The text token sequence is: `tok_1 tok_2 tok_3 ... tok_4`. + - We get the image hidden states for the image through the vision encoder and that hidden state, after a pixel shuffle operation, is then projected into the text embedding space. + We thus have a sequence of image hidden states of size (1, image_seq_len, hidden_dim), where 1 is for batch_size of 1 image and hidden_dim is the hidden_dim of the LM transformer. + - The merging happens so that we obtain the following sequence: `vector_tok_1 vector_tok_2 vector_tok_3 vector_fake_tok_around_image {sequence of image_seq_len image hidden states} vector_fake_toke_around_image vector_tok_4`. That sequence is fed to the LM. + - To fit the format of that sequence, `input_ids`, `inputs_embeds`, `attention_mask` are all 3 adapted to insert the image hidden states. + """ + _, patch_size, _ = image_hidden_states.shape + + if input_ids is None: + image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + image_mask = image_mask[..., 0] # slice off the hidden dim + else: + image_mask = input_ids == self.config.image_token_id + + num_image_tokens = image_mask.sum(dim=1) + torch_compilable_check( + torch.all(num_image_tokens % patch_size == 0), + "At least one sample has tokens not divisible by patch_size.", + ) + blocks_per_sample = num_image_tokens // patch_size + + offsets = torch.nn.functional.pad(blocks_per_sample.cumsum(dim=0), (1, 0), value=0) + block_offset = offsets[:-1] + row_cum = image_mask.cumsum(dim=-1) + chunk_idx = (row_cum - 1) // patch_size + local_idx = (row_cum - 1) % patch_size + block_idx = block_offset.unsqueeze(1) + chunk_idx + + image_embeds = torch.zeros_like(inputs_embeds) + image_embeds[image_mask] = image_hidden_states[block_idx[image_mask], local_idx[image_mask], :] + + merged_embeds = torch.where(image_mask.unsqueeze(-1), image_embeds, inputs_embeds) + return merged_embeds + + @can_return_tuple + @auto_docstring( + custom_intro="Encodes images into continuous embeddings that can be forwarded to the language model." + ) + def get_image_features( + self, + pixel_values: torch.FloatTensor, + pixel_attention_mask: 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. + pixel_attention_mask (`torch.LongTensor`, *optional*): + The attention mask indicating padded regions in the image. + """ + batch_size, num_images, num_channels, height, width = pixel_values.shape + pixel_values = pixel_values.to(dtype=self.dtype) # fp16 compatibility + pixel_values = pixel_values.view(batch_size * num_images, *pixel_values.shape[2:]) + + # Remove padding images - padding images are full 0. + nb_values_per_image = pixel_values.shape[1:].numel() + real_images_inds = (pixel_values == 0.0).sum(dim=(-1, -2, -3)) != nb_values_per_image + + # If no images, leave one empty image. + real_images_inds[0] |= ~torch.any(real_images_inds) + + pixel_values = pixel_values[real_images_inds].contiguous() + # Handle the vision attention mask + if pixel_attention_mask is None: + pixel_attention_mask = torch.ones( + size=[pixel_values.shape[i] for i in (0, 2, 3)], + dtype=torch.bool, + device=pixel_values.device, + ) + else: + # Remove padding images from the mask + pixel_attention_mask = pixel_attention_mask.view(batch_size * num_images, *pixel_attention_mask.shape[2:]) + pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous() + patch_size = self.config.vision_config.patch_size + patches_subgrid = pixel_attention_mask.unfold(dimension=1, size=patch_size, step=patch_size) + patches_subgrid = patches_subgrid.unfold(dimension=2, size=patch_size, step=patch_size) + patch_attention_mask = (patches_subgrid.sum(dim=(-1, -2)) > 0).bool() + + # Get sequence from the vision encoder + image_outputs = self.vision_model( + pixel_values=pixel_values, patch_attention_mask=patch_attention_mask, return_dict=True, **kwargs + ) + image_hidden_states = image_outputs.last_hidden_state + + # Modality projection & resampling + image_features = self.connector(image_hidden_states) + image_outputs.pooler_output = image_features + + return image_outputs + + @can_return_tuple + @auto_docstring( + custom_intro=""" + Inputs fed to the model can have an arbitrary number of images. To account for this, pixel_values fed to + the model have image padding -> (batch_size, max_num_images, 3, max_heights, max_widths) where + max_num_images is the maximum number of images among the batch_size samples in the batch. + Padding images are not needed beyond padding the pixel_values at the entrance of the model. + For efficiency, we only pass through the vision_model's forward the real images by + discarding the padding images i.e. pixel_values of size (image_batch_size, 3, height, width) where + image_batch_size would be 7 when num_images_per_sample=[1, 3, 1, 2] and max_num_images would be 3. + """ + ) + @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.FloatTensor | None = None, + pixel_attention_mask: torch.BoolTensor | None = None, + image_hidden_states: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple | SmolVLMBaseModelOutputWithPast: + r""" + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): + Mask to avoid performing attention on padding pixel indices. + image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The hidden states of the image encoder after modality projection. + """ + if self.training and self.text_model.gradient_checkpointing and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + if input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either 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.text_model.get_input_embeddings()(input_ids).to(input_ids.device) + + if pixel_values is not None and image_hidden_states is not None: + raise ValueError("You cannot specify both pixel_values and image_hidden_states at the same time") + + if pixel_values is not None: + image_hidden_states = self.get_image_features( + pixel_values, pixel_attention_mask, return_dict=True + ).pooler_output + image_hidden_states = image_hidden_states.to(inputs_embeds.device) + elif image_hidden_states is not None: + image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=inputs_embeds.device) + + if image_hidden_states is not None: + inputs_embeds = self.inputs_merger( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + image_hidden_states=image_hidden_states, + ) + + outputs = self.text_model( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + return SmolVLMBaseModelOutputWithPast( + 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_hidden_states, + ) + + +@dataclass +@auto_docstring( + custom_intro=""" + Base class for Idefics causal language model (or autoregressive) outputs. + """ +) +class SmolVLMCausalLMOutputWithPast(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 + """ + + 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 + + +@auto_docstring( + custom_intro=""" + The SmolVLM Model with a language modeling head. It is made up a SigLIP vision encoder, with a language modeling head on top. + """ +) +class SmolVLMForConditionalGeneration(SmolVLMPreTrainedModel, GenerationMixin): + _tied_weights_keys = {"lm_head.weight": "model.text_model.embed_tokens.weight"} + + def __init__(self, config): + super().__init__(config) + self.model = SmolVLMModel(config) + self.image_token_id = self.config.image_token_id + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.vocab_size = config.text_config.vocab_size + self.model.text_model.generation_config = GenerationConfig.from_model_config(config) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self): + return self.model.text_model.get_input_embeddings() + + def set_input_embeddings(self, value): + self.model.text_model.set_input_embeddings(value) + + @auto_docstring + def get_image_features( + self, + pixel_values: torch.FloatTensor, + pixel_attention_mask: 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. + pixel_attention_mask (`torch.LongTensor`, *optional*): + The attention mask indicating padded regions in the image. + """ + return self.model.get_image_features( + pixel_values=pixel_values, pixel_attention_mask=pixel_attention_mask, **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, + pixel_values: torch.FloatTensor | None = None, + pixel_attention_mask: torch.BoolTensor | None = None, + image_hidden_states: 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 | SmolVLMCausalLMOutputWithPast: + r""" + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): + Mask to avoid performing attention on padding pixel indices. + image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The hidden states of the image encoder after modality projection. + 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 `model.image_token_id`. Tokens with indices set to `model.image_token_id` are + ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Example: + + ```python + >>> import httpx + >>> from io import BytesIO + >>> import torch + >>> from PIL import Image + >>> from io import BytesIO + + >>> from transformers import AutoProcessor, AutoModelForImageTextToText + >>> from transformers.image_utils import load_image + + >>> # Note that passing the image urls (instead of the actual pil images) to the processor is also possible + >>> image1 = load_image("https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg") + >>> image2 = load_image("https://cdn.britannica.com/59/94459-050-DBA42467/Skyline-Chicago.jpg") + >>> image3 = load_image("https://cdn.britannica.com/68/170868-050-8DDE8263/Golden-Gate-Bridge-San-Francisco.jpg") + + >>> processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct") + >>> model = AutoModelForImageTextToText.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct", dtype=torch.bfloat16, device_map="auto") + + >>> # Create inputs + >>> messages = [ + ... { + ... "role": "user", + ... "content": [ + ... {"type": "video", "path": path/to/video}, + ... {"type": "text", "text": "What is happening in this video?"}, + ... ] + ... } + ... ] + + >>> inputs = processor.apply_chat_template([messages], add_generation_prompt=True) + + >>> # Generate + >>> generated_ids = model.generate(**inputs, max_new_tokens=256) + >>> generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True) + + >>> print(generated_texts) + ```""" + # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_attention_mask=pixel_attention_mask, + image_hidden_states=image_hidden_states, + use_cache=use_cache, + return_dict=True, + **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 SmolVLMCausalLMOutputWithPast( + 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, + attention_mask=None, + inputs_embeds=None, + pixel_values=None, + pixel_attention_mask=None, + image_hidden_states=None, + logits_to_keep=None, + is_first_iteration=False, + use_cache=False, + **kwargs, + ): + # Overwritten -- there are mutually exclusive inputs (if the logic to make `image_hidden_states` take + # precedence is moved to the model, we can remove this fn) + + model_inputs = super().prepare_inputs_for_generation( + input_ids, + past_key_values=past_key_values, + attention_mask=attention_mask, + inputs_embeds=inputs_embeds, + pixel_values=pixel_values, + pixel_attention_mask=pixel_attention_mask, + image_hidden_states=image_hidden_states, + logits_to_keep=logits_to_keep, + is_first_iteration=is_first_iteration, + use_cache=use_cache, + **kwargs, + ) + + if image_hidden_states is not None or (use_cache and not is_first_iteration): + model_inputs["pixel_values"] = None + model_inputs["pixel_attention_mask"] = None + + return model_inputs + + +__all__ = ["SmolVLMForConditionalGeneration", "SmolVLMPreTrainedModel", "SmolVLMModel", "SmolVLMVisionTransformer"] diff --git a/third_party/transformers/src/transformers/models/smolvlm/modular_smolvlm.py b/third_party/transformers/src/transformers/models/smolvlm/modular_smolvlm.py new file mode 100644 index 0000000000000000000000000000000000000000..9c572cc9d8774fa02e65ae63a8cf369f90e4c87a --- /dev/null +++ b/third_party/transformers/src/transformers/models/smolvlm/modular_smolvlm.py @@ -0,0 +1,370 @@ +# Copyright 2025 the HuggingFace Inc. team. All rights reserved. +# Written by Orr Zohar +# +# 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 ...cache_utils import Cache, DynamicCache +from ...generation import GenerationConfig +from ...modeling_flash_attention_utils import FlashAttentionKwargs +from ...modeling_outputs import BaseModelOutputWithPooling +from ...processing_utils import ImagesKwargs, Unpack +from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging, torch_compilable_check +from ..idefics3.configuration_idefics3 import Idefics3Config, Idefics3VisionConfig +from ..idefics3.image_processing_idefics3 import Idefics3ImageProcessor +from ..idefics3.image_processing_pil_idefics3 import Idefics3ImageProcessorPil +from ..idefics3.modeling_idefics3 import ( + Idefics3BaseModelOutputWithPast, + Idefics3ForConditionalGeneration, + Idefics3Model, + Idefics3PreTrainedModel, + Idefics3VisionTransformer, +) + + +logger = logging.get_logger(__name__) + + +@auto_docstring(checkpoint="HuggingFaceTB/SmolVLM2-2.2B-Instruct") +@strict +class SmolVLMVisionConfig(Idefics3VisionConfig): + r""" + Example: + + ```python + >>> from transformers.models.smolvlm.modeling_smolvlm import SmolVLMVisionTransformer + >>> from transformers.models.smolvlm.configuration_smolvlm import SmolVLMVisionConfig + + >>> # Initializing a SmolVLMVisionConfig with google/siglip-so400m-patch14-384 style configuration + >>> configuration = SmolVLMVisionConfig() + + >>> # Initializing a SmolVLMVisionTransformer (with random weights) from the google/siglip-so400m-patch14-384 style configuration + >>> model = SmolVLMVisionTransformer(configuration) + + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "smolvlm_vision" + + +class SmolVLMPreTrainedModel(Idefics3PreTrainedModel): + pass + + +class SmolVLMVisionTransformer(Idefics3VisionTransformer): + pass + + +@auto_docstring(checkpoint="HuggingFaceTB/SmolVLM2-2.2B-Instruct") +@strict +class SmolVLMConfig(Idefics3Config): + r""" + scale_factor (`int`, *optional*, defaults to 2): + The scale factor for the image encoder. + + Example: + ```python + >>> from transformers import SmolVLMModel, SmolVLMConfig + >>> # Initializing configuration + >>> configuration = SmolVLMConfig() + >>> # Initializing a model from the configuration + >>> model = SmolVLMModel(configuration) + >>> # Accessing the model configuration + >>> configuration = model.config + ```""" + + model_type = "smolvlm" + + +class SmolVLMImageProcessorKwargs(ImagesKwargs, total=False): + """ + do_image_splitting (`bool`, *optional*, defaults to `True`): + Whether to split the image into sub-images concatenated with the original image. They are split into patches + such that each patch has a size of `max_image_size["height"]` x `max_image_size["width"]`. + max_image_size (`Dict`, *optional*, defaults to `{"longest_edge": 364}`): + Maximum resolution of the patches of images accepted by the model. This is a dictionary containing the key "longest_edge". + return_row_col_info (`bool`, *optional*, defaults to `False`): + Whether to return the row and column information of the images. + """ + + do_image_splitting: bool + max_image_size: dict[str, int] + return_row_col_info: bool + + +class SmolVLMImageProcessor(Idefics3ImageProcessor): + pass + + +class SmolVLMImageProcessorPil(Idefics3ImageProcessorPil): + pass + + +class SmolVLMBaseModelOutputWithPast(Idefics3BaseModelOutputWithPast): + pass + + +class SmolVLMModel(Idefics3Model): + """ + A subclass of Idefics3Model. We do *not* remove or block the call to inputs_merger + in forward. Instead, we override inputs_merger here with custom logic. + """ + + def inputs_merger( + self, input_ids: torch.LongTensor, inputs_embeds: torch.Tensor, image_hidden_states: torch.Tensor + ): + _, patch_size, _ = image_hidden_states.shape + + if input_ids is None: + image_mask = inputs_embeds == self.get_input_embeddings()( + torch.tensor(self.config.image_token_id, dtype=torch.long, device=inputs_embeds.device) + ) + image_mask = image_mask[..., 0] # slice off the hidden dim + else: + image_mask = input_ids == self.config.image_token_id + + num_image_tokens = image_mask.sum(dim=1) + torch_compilable_check( + torch.all(num_image_tokens % patch_size == 0), + "At least one sample has tokens not divisible by patch_size.", + ) + blocks_per_sample = num_image_tokens // patch_size + + offsets = torch.nn.functional.pad(blocks_per_sample.cumsum(dim=0), (1, 0), value=0) + block_offset = offsets[:-1] + row_cum = image_mask.cumsum(dim=-1) + chunk_idx = (row_cum - 1) // patch_size + local_idx = (row_cum - 1) % patch_size + block_idx = block_offset.unsqueeze(1) + chunk_idx + + image_embeds = torch.zeros_like(inputs_embeds) + image_embeds[image_mask] = image_hidden_states[block_idx[image_mask], local_idx[image_mask], :] + + merged_embeds = torch.where(image_mask.unsqueeze(-1), image_embeds, inputs_embeds) + return merged_embeds + + @can_return_tuple + @auto_docstring( + custom_intro="Encodes images into continuous embeddings that can be forwarded to the language model." + ) + def get_image_features( + self, + pixel_values: torch.FloatTensor, + pixel_attention_mask: 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. + pixel_attention_mask (`torch.LongTensor`, *optional*): + The attention mask indicating padded regions in the image. + """ + batch_size, num_images, num_channels, height, width = pixel_values.shape + pixel_values = pixel_values.to(dtype=self.dtype) # fp16 compatibility + pixel_values = pixel_values.view(batch_size * num_images, *pixel_values.shape[2:]) + + # Remove padding images - padding images are full 0. + nb_values_per_image = pixel_values.shape[1:].numel() + real_images_inds = (pixel_values == 0.0).sum(dim=(-1, -2, -3)) != nb_values_per_image + + # If no images, leave one empty image. + real_images_inds[0] |= ~torch.any(real_images_inds) + + pixel_values = pixel_values[real_images_inds].contiguous() + # Handle the vision attention mask + if pixel_attention_mask is None: + pixel_attention_mask = torch.ones( + size=[pixel_values.shape[i] for i in (0, 2, 3)], + dtype=torch.bool, + device=pixel_values.device, + ) + else: + # Remove padding images from the mask + pixel_attention_mask = pixel_attention_mask.view(batch_size * num_images, *pixel_attention_mask.shape[2:]) + pixel_attention_mask = pixel_attention_mask[real_images_inds].contiguous() + patch_size = self.config.vision_config.patch_size + patches_subgrid = pixel_attention_mask.unfold(dimension=1, size=patch_size, step=patch_size) + patches_subgrid = patches_subgrid.unfold(dimension=2, size=patch_size, step=patch_size) + patch_attention_mask = (patches_subgrid.sum(dim=(-1, -2)) > 0).bool() + + # Get sequence from the vision encoder + image_outputs = self.vision_model( + pixel_values=pixel_values, patch_attention_mask=patch_attention_mask, return_dict=True, **kwargs + ) + image_hidden_states = image_outputs.last_hidden_state + + # Modality projection & resampling + image_features = self.connector(image_hidden_states) + image_outputs.pooler_output = image_features + + return image_outputs + + @can_return_tuple + @auto_docstring( + custom_intro=""" + Inputs fed to the model can have an arbitrary number of images. To account for this, pixel_values fed to + the model have image padding -> (batch_size, max_num_images, 3, max_heights, max_widths) where + max_num_images is the maximum number of images among the batch_size samples in the batch. + Padding images are not needed beyond padding the pixel_values at the entrance of the model. + For efficiency, we only pass through the vision_model's forward the real images by + discarding the padding images i.e. pixel_values of size (image_batch_size, 3, height, width) where + image_batch_size would be 7 when num_images_per_sample=[1, 3, 1, 2] and max_num_images would be 3. + """ + ) + @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.FloatTensor | None = None, + pixel_attention_mask: torch.BoolTensor | None = None, + image_hidden_states: torch.FloatTensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple | SmolVLMBaseModelOutputWithPast: + if self.training and self.text_model.gradient_checkpointing and use_cache: + logger.warning_once( + "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..." + ) + use_cache = False + + if input_ids is not None: + batch_size, seq_length = input_ids.shape + elif inputs_embeds is not None: + batch_size, seq_length, _ = inputs_embeds.shape + else: + raise ValueError("You have to specify either 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.text_model.get_input_embeddings()(input_ids).to(input_ids.device) + + if pixel_values is not None and image_hidden_states is not None: + raise ValueError("You cannot specify both pixel_values and image_hidden_states at the same time") + + if pixel_values is not None: + image_hidden_states = self.get_image_features( + pixel_values, pixel_attention_mask, return_dict=True + ).pooler_output + image_hidden_states = image_hidden_states.to(inputs_embeds.device) + elif image_hidden_states is not None: + image_hidden_states = image_hidden_states.to(dtype=self.dtype, device=inputs_embeds.device) + + if image_hidden_states is not None: + inputs_embeds = self.inputs_merger( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + image_hidden_states=image_hidden_states, + ) + + outputs = self.text_model( + inputs_embeds=inputs_embeds, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + **kwargs, + ) + + return SmolVLMBaseModelOutputWithPast( + 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_hidden_states, + ) + + +class SmolVLMForConditionalGeneration(Idefics3ForConditionalGeneration): + _tied_weights_keys = {"lm_head.weight": "model.text_model.embed_tokens.weight"} + + def __init__(self, config): + super().__init__(config) + self.model = SmolVLMModel(config) + self.model.text_model.generation_config = GenerationConfig.from_model_config(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.post_init() + + def forward(self, **super_kwargs): + r""" + pixel_attention_mask (`torch.Tensor` of shape `(batch_size, image_size, image_size)`, *optional*): + Mask to avoid performing attention on padding pixel indices. + image_hidden_states (`torch.FloatTensor` of shape `(batch_size, num_channels, image_size, image_size)`): + The hidden states of the image encoder after modality projection. + 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 `model.image_token_id`. Tokens with indices set to `model.image_token_id` are + ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. + + Example: + + ```python + >>> import httpx + >>> from io import BytesIO + >>> import torch + >>> from PIL import Image + >>> from io import BytesIO + + >>> from transformers import AutoProcessor, AutoModelForImageTextToText + >>> from transformers.image_utils import load_image + + >>> # Note that passing the image urls (instead of the actual pil images) to the processor is also possible + >>> image1 = load_image("https://cdn.britannica.com/61/93061-050-99147DCE/Statue-of-Liberty-Island-New-York-Bay.jpg") + >>> image2 = load_image("https://cdn.britannica.com/59/94459-050-DBA42467/Skyline-Chicago.jpg") + >>> image3 = load_image("https://cdn.britannica.com/68/170868-050-8DDE8263/Golden-Gate-Bridge-San-Francisco.jpg") + + >>> processor = AutoProcessor.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct") + >>> model = AutoModelForImageTextToText.from_pretrained("HuggingFaceTB/SmolVLM2-2.2B-Instruct", dtype=torch.bfloat16, device_map="auto") + + >>> # Create inputs + >>> messages = [ + ... { + ... "role": "user", + ... "content": [ + ... {"type": "video", "path": path/to/video}, + ... {"type": "text", "text": "What is happening in this video?"}, + ... ] + ... } + ... ] + + >>> inputs = processor.apply_chat_template([messages], add_generation_prompt=True) + + >>> # Generate + >>> generated_ids = model.generate(**inputs, max_new_tokens=256) + >>> generated_texts = processor.batch_decode(generated_ids, skip_special_tokens=True) + + >>> print(generated_texts) + ```""" + super().forward(**super_kwargs) + + +__all__ = [ + "SmolVLMVisionConfig", + "SmolVLMConfig", + "SmolVLMImageProcessor", + "SmolVLMImageProcessorPil", + "SmolVLMForConditionalGeneration", + "SmolVLMPreTrainedModel", + "SmolVLMModel", + "SmolVLMVisionTransformer", +] diff --git a/third_party/transformers/src/transformers/models/smolvlm/processing_smolvlm.py b/third_party/transformers/src/transformers/models/smolvlm/processing_smolvlm.py new file mode 100644 index 0000000000000000000000000000000000000000..69eb7d6e5a617fe2495da4fbf064c48e57faf458 --- /dev/null +++ b/third_party/transformers/src/transformers/models/smolvlm/processing_smolvlm.py @@ -0,0 +1,353 @@ +# 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 SmolVLM. +""" + +from datetime import timedelta +from typing import TYPE_CHECKING, Union + +from ...feature_extraction_utils import BatchFeature +from ...image_utils import ImageInput, make_nested_list_of_images +from ...processing_utils import ProcessingKwargs, ProcessorMixin, Unpack +from ...tokenization_utils_base import BatchEncoding, TextInput +from ...utils import auto_docstring, is_num2words_available, logging +from ...video_utils import VideoInput + + +# Adapted from transformers.models.smolvlm.video_processing_smolvlm.DEFAULT_VIDEO_INTRO +DEFAULT_VIDEO_INTRO = ( + "You are provided the following series of {frame_count} frames from a {video_duration} [H:MM:SS] video.\n" +) +# Adapted from transformers.models.smolvlm.video_processing_smolvlm.DEFAULT_MEDIA_OUTTRO +DEFAULT_MEDIA_OUTTRO = "\n\n" +# Adapted from transformers.models.smolvlm.video_processing_smolvlm.FRAME_TIMESTAMP_MESSAGE +FRAME_TIMESTAMP_MESSAGE = "\nFrame from {timestamp}:" + +if TYPE_CHECKING: + from ...tokenization_utils_base import PreTokenizedInput + +logger = logging.get_logger(__name__) + + +if is_num2words_available(): + from num2words import num2words +else: + num2words = None + + +# The correct chat template to be used for videos after #38105 +DEFAULT_CHAT_TEMPLATE = "<|im_start|>{% for message in messages %}{{message['role'] | capitalize}}{% if message['content'][0]['type'] == 'image' %}{{':'}}{% else %}{{': '}}{% endif %}{% for line in message['content'] %}{% if line['type'] == 'text' %}{{line['text']}}{% elif line['type'] == 'image' %}{{ '' }}{% elif line['type'] == 'video' %}{{ '